Diagonal Traverse II

medium heap priority queue matrix sorting

Problem

Given a (possibly jagged) 2D integer array nums, return all of its elements in diagonal order. Cells on the same anti-diagonal share the same value of i + j; smaller diagonals come first, and within a diagonal cells are listed from the bottom-left up to the top-right (decreasing row index order in the output reads bottom row first).

Inputnums = [[1,2,3],[4,5],[6]]

Output[1,4,2,6,5,3]

Diagonal 0: [1]; diagonal 1: [4,2]; diagonal 2: [6,5,3]. Within each diagonal the lower rows come first.

Inputnums = [[1,2,3],[4,5,6],[7,8,9]]

Output[1,4,2,7,5,3,8,6,9]

A full 3×3 grid traversed by anti-diagonals.

nums (rows separated by ;, values by ,)

def findDiagonalOrder(nums):
    import heapq
    heap = []
    for i in range(len(nums)):
        for j in range(len(nums[i])):
            heapq.heappush(heap, (i + j, -i, nums[i][j]))
    res = []
    while heap:
        diag, neg_i, val = heapq.heappop(heap)
        res.append(val)
    return res

function findDiagonalOrder(nums) {
  const heap = [];                 // [diag, -i, val] tuples
  for (let i = 0; i < nums.length; i++)
    for (let j = 0; j < nums[i].length; j++)
      heap.push([i + j, -i, nums[i][j]]);
  heap.sort((a, b) => a[0] - b[0] || a[1] - b[1]);
  const res = [];
  for (const t of heap)
    res.push(t[2]);
  return res;
}

int[] findDiagonalOrder(List<List<Integer>> nums) {
    PriorityQueue<int[]> heap = new PriorityQueue<>(
        (a, b) -> a[0] != b[0] ? a[0] - b[0] : a[1] - b[1]);
    int total = 0;
    for (int i = 0; i < nums.size(); i++)
        for (int j = 0; j < nums.get(i).size(); j++) {
            heap.offer(new int[]{i + j, -i, nums.get(i).get(j)});
            total++;
        }
    int[] res = new int[total];
    for (int k = 0; k < total; k++)
        res[k] = heap.poll()[2];
    return res;
}

vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
    priority_queue<array<int,3>, vector<array<int,3>>,
                   greater<>> heap;
    for (int i = 0; i < (int)nums.size(); i++)
        for (int j = 0; j < (int)nums[i].size(); j++)
            heap.push({i + j, -i, nums[i][j]});
    vector<int> res;
    while (!heap.empty()) {
        res.push_back(heap.top()[2]);
        heap.pop();
    }
    return res;
}

Explanation

Every cell nums[i][j] lies on the anti-diagonal numbered d = i + j. Diagonal 0 is just the top-left corner; each step out adds one more diagonal. The required output lists the diagonals in increasing order of d, and inside one diagonal it lists cells from the lowest row upward.

That ordering is exactly a sort by the key (i + j, -i): the primary key i + j groups and orders the diagonals, and the secondary key -i makes larger row indices come first within a diagonal. A min-heap / priority queue keyed by this tuple emits the cells in precisely the right sequence.

So we push every cell as (i + j, -i, value) into the heap, then repeatedly pop the smallest tuple and append its value to the answer. The heap handles the grouping and the within-diagonal ordering for us, no matter how ragged the rows are.

Because each cell is pushed and popped once, the work is dominated by the heap operations. With N total cells that is O(N log N) time and O(N) space. (A bucket-by-i+j approach can reach O(N), but the heap version is the clearest expression of the ordering rule.)

For nums = [[1,2,3],[4,5],[6]] the heap pops 1 (d=0), then 4 then 2 (d=1, row 1 before row 0), then 6, 5, 3 (d=2, rows 2,1,0), giving [1,4,2,6,5,3].

Time: O(N log N) Space: O(N)