1 0  1  2  3
2 4  5  6  7
3 8  9 10 11
412 13 14 15

1def zigzag_order(matrix):
2    if not matrix:
3        return []
4    rows, cols = len(matrix), len(matrix[0])
5    result = []
6
7    for diag in range(rows + cols - 1):
8        if diag % 2 == 0:
9            # Even diagonal: move upward (row decreases, col increases)
10            r = min(diag, rows - 1)
11            c = diag - r
12            while r >= 0 and c < cols:
13                result.append(matrix[r][c])
14                r -= 1
15                c += 1
16        else:
17            # Odd diagonal: move downward (row increases, col decreases)
18            c = min(diag, cols - 1)
19            r = diag - c
20            while c >= 0 and r < rows:
21                result.append(matrix[r][c])
22                r += 1
23                c -= 1
24
25    return result
26
27matrix = [
28    [0,  1,  2,  3],
29    [4,  5,  6,  7],
30    [8,  9, 10, 11],
31    [12, 13, 14, 15]
32]
33print(zigzag_order(matrix))
34# [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15]

1def zigzag_ith(n, m, i):
2    """Find the i-th element (0-indexed) in zigzag order of an n x m matrix."""
3    count = 0
4
5    for diag in range(n + m - 1):
6        # Number of elements in this diagonal
7        if diag < min(n, m):
8            diag_len = diag + 1
9        elif diag < max(n, m):
10            diag_len = min(n, m)
11        else:
12            diag_len = n + m - 1 - diag
13
14        if count + diag_len > i:
15            # The i-th element is in this diagonal
16            pos = i - count  # position within diagonal
17
18            if diag % 2 == 0:
19                # Upward: start at (min(diag, n-1), diag - min(diag, n-1))
20                r = min(diag, n - 1) - pos
21                c = diag - min(diag, n - 1) + pos
22            else:
23                # Downward: start at (diag - min(diag, m-1), min(diag, m-1))
24                r = diag - min(diag, m - 1) + pos
25                c = min(diag, m - 1) - pos
26
27            return (r, c)
28
29        count += diag_len
30
31    return None  # i out of bounds
32
33# Find the 7th element (0-indexed) in a 4x4 matrix
34r, c = zigzag_ith(4, 4, 7)
35print(f"Position: ({r}, {c}), Value: {matrix[r][c]}")
36# Position: (1, 2), Value: 6

1def zigzag_ith_square(n, i):
2    """O(1) lookup for i-th element in zigzag order of an n x n matrix."""
3    import math
4
5    # Find which diagonal: diag d has d+1 elements (for d < n)
6    # Sum of first d diagonals = d*(d+1)/2
7    # Solve d*(d+1)/2 <= i for d
8    d = int((-1 + math.sqrt(1 + 8 * i)) / 2)
9
10    # Adjust if we overshot
11    while (d + 1) * (d + 2) // 2 <= i:
12        d += 1
13
14    # Position within diagonal
15    pos = i - d * (d + 1) // 2
16
17    if d % 2 == 0:
18        return (d - pos, pos)      # Upward
19    else:
20        return (pos, d - pos)      # Downward
21
22# Verify against full zigzag
23for idx in range(16):
24    r, c = zigzag_ith_square(4, idx)
25    print(f"i={idx}: ({r},{c}) = {matrix[r][c]}")

1# Standard JPEG zigzag table for 8x8 block
2JPEG_ZIGZAG = [
3    0,  1,  8, 16,  9,  2,  3, 10,
4   17, 24, 32, 25, 18, 11,  4,  5,
5   12, 19, 26, 33, 40, 48, 41, 34,
6   27, 20, 13,  6,  7, 14, 21, 28,
7   35, 42, 49, 56, 57, 50, 43, 36,
8   29, 22, 15, 23, 30, 37, 44, 51,
9   58, 59, 52, 45, 38, 31, 39, 46,
10   53, 60, 61, 54, 47, 55, 62, 63
11]
12
13def zigzag_to_block(coefficients):
14    """Reorder 64 zigzag-ordered DCT coefficients into 8x8 block."""
15    block = [[0]*8 for _ in range(8)]
16    for zigzag_idx, coeff in enumerate(coefficients):
17        linear_pos = JPEG_ZIGZAG[zigzag_idx]
18        block[linear_pos // 8][linear_pos % 8] = coeff
19    return block

1def position_to_zigzag_index(n, m, r, c):
2    """Find zigzag index for position (r, c) in an n x m matrix."""
3    diag = r + c
4    # Count elements in all previous diagonals
5    count = 0
6    for d in range(diag):
7        if d < min(n, m):
8            count += d + 1
9        elif d < max(n, m):
10            count += min(n, m)
11        else:
12            count += n + m - 1 - d
13
14    # Position within this diagonal
15    if diag % 2 == 0:
16        pos = c  # Upward diagonal, position = column
17    else:
18        pos = r  # Downward diagonal, position = row
19
20    # Adjust for diagonals beyond the first row/column
21    if diag >= n and diag % 2 == 0:
22        pos = c - (diag - n + 1)
23    if diag >= m and diag % 2 == 1:
24        pos = r - (diag - m + 1)
25
26    return count + pos
27
28print(position_to_zigzag_index(4, 4, 1, 2))  # 7 (value 6 is at zigzag index 7)