How to find ith item in zigzag ordering?

Question

A question last week defined the zig zag ordering on an n by m matrix and asked how to list the elements in that order.

My question is how to quickly find the ith item in the zigzag ordering? That is, without traversing the matrix (for large n and m that's much too slow).

For example with n=m=8 as in the picture and (x, y) describing (row, column)

f(0) = (0, 0)
f(1) = (0, 1)
f(2) = (1, 0)
f(3) = (2, 0)
f(4) = (1, 1)
...
f(63) = (7, 7)

Specific question: what is the ten billionth (1e10) item in the zigzag ordering of a million by million matrix?

Answer 1

1. Let's assume that the desired element is located in the upper half of the matrix. The length of the diagonals are 1, 2, 3 ..., n.

2. Let's find the desired diagonal. It satisfies the following property: sum(1, 2 ..., k) >= pos but sum(1, 2, ..., k - 1) < pos. The sum of 1, 2, ..., k is k * (k + 1) / 2. So we just need to find the smallest integer k such that k * (k + 1) / 2 >= pos. We can either use a binary search or solve this quadratic inequality explicitly.

3. When we know the k, we just need to find the pos - (k - 1) * k / 2 element of this diagonal. We know where it starts and where we should move(up or down, depending on the parity of k), so we can find the desired cell using a simple formula.

This solution has an O(1) or an O(log n) time complexity(it depends on whether we use a binary search or solve the inequation explicitly in step 2).

If the desired element is located in the lower half of the matrix, we can solve this problem for a pos' = n * n - pos + 1 and then use symmetry to get the solution to the original problem.

I used 1-based indexing in this solution, using 0-based indexing might require adding +1 or -1 somewhere, but the idea of the solution is the same.

If the matrix is rectangular, not square, we need to consider the fact the length of diagonals look this way: 1, 2, 3, ..., m, m, m, .., m, m - 1, ..., 1(if m <= n) when we search for the k, so the sum becomes something like k * (k + 1) / 2 if k <= m and k * (k + 1) / 2 + m * (k - m) otherwise.

Answer 2

import math, random

def naive(n, m, ord, swap = False):
    dx = 1
    dy = -1
    if swap:
        dx, dy = dy, dx
    cur = [0, 0]
    for i in range(ord):
        cur[0] += dy
        cur[1] += dx
        if cur[0] < 0 or cur[1] < 0 or cur[0] >= n or cur[1] >= m:
            dx, dy = dy, dx

        if cur[0] >= n:
            cur[0] = n - 1
            cur[1] += 2
        if cur[1] >= m:
            cur[1] = m - 1
            cur[0] += 2
        if cur[0] < 0: cur[0] = 0
        if cur[1] < 0: cur[1] = 0


    return cur

def fast(n, m, ord, swap = False):
    if n < m:
        x, y = fast(m, n, ord, not swap)
        return [y, x]

    alt = n * m - ord - 1
    if alt < ord:
        x, y = fast(n, m, alt, swap if (n + m) % 2 == 0 else not swap)
        return [n - x - 1, m - y - 1]

    if ord < (m * (m + 1) / 2):
        diag = int((-1 + math.sqrt(1 + 8 * ord)) / 2)
        parity = (diag + (0 if swap else 1)) % 2
        within = ord - (diag * (diag + 1) / 2)
        if parity: return [diag - within, within]
        else: return [within, diag - within]
    else:
        ord -= (m * (m + 1) / 2)
        diag = int(ord / m)
        within = ord - diag * m
        diag += m
        parity = (diag + (0 if swap else 1)) % 2
        if not parity:
            within = m - within - 1
        return [diag - within, within]

if __name__ == "__main__":
    for i in range(1000):
        n = random.randint(3, 100)
        m = random.randint(3, 100)
        ord = random.randint(0, n * m - 1)
        swap = random.randint(0, 99) < 50
        na = naive(n, m, ord, swap)
        fa = fast(n, m, ord, swap)
        assert na == fa, "(%d, %d, %d, %s) ==> (%s), (%s)" % (n, m, ord, swap, na, fa)

    print fast(1000000, 1000000, 9999999999, False)
    print fast(1000000, 1000000, 10000000000, False)

So the 10-billionth element (the one with ordinal 9999999999), and the 10-billion-first element (the one with ordinal 10^10) are:

[20331, 121089]
[20330, 121090]