index to coordinates in diagonal (zigzag) traverse

Question

So I have a 4x4 matrix like so

 |0 1 2 3
-+-------
0|0 1 3 6
1|2 4 7 a
2|5 8 b d
3|9 c e f

and I am traversing it in the order specified by the hexadecimal characters in it. so start at (0, 0), then (1, 0), (0, 1), (2, 0), (1, 1), (0, 2)...

So here is the code:

def diagonal(n):
    for a in range(n):
        for b in range(a + 1):
            yield a - b, b
    for a in range(n - 1):
        for b in range(n - a - 1):
            yield n - b - 1, b + 1 + a

iterating through this gives

for x, y in diagonal(4):
    print((x, y))

# (0, 0)
# (1, 0)
# (0, 1)
# (2, 0)
# (1, 1)
# (0, 2)
# (3, 0)
# (2, 1)
# (1, 2)
# (0, 3)
# (3, 1)
# (2, 2)
# (1, 3)
# (3, 2)
# (2, 3)
# (3, 3)

which is exactly as I intended. The part I am stuck on is trying to make a function where I give it the index, and it gives me the coordinates. So for my 4x4 matrix,(this is still in hexadecimal)

0 -> (0, 0)
1 -> (1, 0)
2 -> (0, 1)
3 -> (2, 0)
4 -> (1, 1)
5 -> (0, 2)
6 -> (3, 0)
7 -> (2, 1)
8 -> (1, 2)
9 -> (0, 3)
a -> (3, 1)
b -> (2, 2)
c -> (1, 3)
d -> (3, 2)
e -> (2, 3)
f -> (3, 3)

I have intentions to move on to variable sized square matrices so I cannot just hard code these values into a dict.

I have been tinkering for hours trying to get this to work but I cannot for the life of me get it to work.

This is not homework, just something I'm working on in my spare time, and it's slowly driving me up the wall.

If anything is not clear please don't hesitate to ask.

Thanks in advance.

Edit: I presume someone will comment about this post Traverse Matrix in Diagonal strips which is similar but as with my first function this only iterates over the coordinates, and I cannot work out the coordinates from an index.

Answer 1

Here is a function that seems to do what you want. An explanation comes after the code.

from math import sqrt

def triangular(n):
    return n * (n + 1) // 2

def coords_from_index(ndx, n=4):
    if ndx < triangular(n):
        basecol = (int(sqrt(8 * ndx + 1)) - 1) // 2
        row = ndx - triangular(basecol)
        col = basecol - row
    else:
        oldcol, oldrow = coords_from_index(n**2 - 1 - ndx, n)
        row = n - 1 - oldrow
        col = n - 1 - oldcol
    return col, row

# Test code
n = 4
for ndx in range(n**2):
    print(hex(ndx)[2:], '->', coords_from_index(ndx, n))

The printout from that test code is:

0 -> (0, 0)
1 -> (1, 0)
2 -> (0, 1)
3 -> (2, 0)
4 -> (1, 1)
5 -> (0, 2)
6 -> (3, 0)
7 -> (2, 1)
8 -> (1, 2)
9 -> (0, 3)
a -> (3, 1)
b -> (2, 2)
c -> (1, 3)
d -> (3, 2)
e -> (2, 3)
f -> (3, 3)

---

Here is a brief explanation of my code.

Like your code to generate the coordinates in order, I treat the upper-left triangle of the square differently from the lower-right triangle. Let's look first at the upper-left triangle, which for the square of size 4 includes the indices 0 through 9.

If you look at the top number in each column, you see that those are the "triangular numbers", which are the sum of consecutive integers starting from 0. So the top row is 0, 0+1, 0+1+2, and 0+1+2+3. The well-known formula for those numbers is

triangular(n) = n * (n + 1) // 2

So I wrote a small routine for that. If you know the triangular number (call it ndx) and want to find n, you can use algebra to solve the quadratic equation and you get

n = (sqrt(8 * ndx + 1) - 1) // 2

If you replace the sqrt with int(sqrt( you get that same result for a triangular number, and you also get the "base" number for any ndx that falls between two triangular numbers. You can then use your index and the "base number" to find the corresponding column and row for the index.

Now, if you look at the lower-right triangle, which includes the indices a through f, you see that there is symmetry with the upper-left triangle. I chose to use that symmetry to calculate the row and column for any of these indices. I could have calculated them more directly, but the approach I used works well.

Note that if you use really large values of n and ndx that the int(sqrt( does not always give the correct answer, due to the vagaries of floating-point arithmetic. But ndx needs to be really large, on the order of 2**50, before that happens. Let me know if you need a more reliable routine to calculate the integer square root.