LeetCode 1424 - Diagonal Traverse II

The problem asks us to traverse a jagged 2D integer array nums diagonally. Specifically, elements are accessed by their

leetcodemediumarraysortingheap-(priority-queue)

LeetCode Problem 1424

Difficulty: 🟡 Medium
Topics: Array, Sorting, Heap (Priority Queue)

Solution

Problem Understanding

The problem asks us to traverse a jagged 2D integer array nums diagonally. Specifically, elements are accessed by their "diagonal level," where each diagonal contains elements with the same sum of row and column indices. Within each diagonal, elements should appear in reverse order of their rows, meaning we start from the bottom-most row of the diagonal and move up.

The input nums is a list of lists of integers. The lists are jagged, meaning each sublist may have a different length. The output should be a single list of integers, representing the diagonal traversal of all elements according to the above rules.

Constraints are significant: the total number of elements across all sublists can be up to 10^5. This rules out any approach that would involve fully flattening or sorting in quadratic time. Each element can be large but the range is not critical for traversal logic.

Edge cases include empty sublists (though nums[i].length >= 1 guarantees at least one element per row), highly jagged arrays (rows of widely differing lengths), and minimal input (nums containing a single row or a single column).

Approaches

The brute-force approach would be to iterate over all possible diagonal levels (sum of indices), collect elements belonging to each diagonal into temporary arrays, reverse each diagonal, and then append them to the result. This works because diagonals are naturally grouped by row + column, but it is cumbersome and may involve many temporary arrays.

The optimal approach uses a hash map (dictionary) where the key is the diagonal level d = row + column. As we iterate over each element, we append it to the list corresponding to its diagonal. After processing all elements, we concatenate the diagonals in order of increasing diagonal level, reversing each diagonal because elements are collected top-to-bottom but must appear bottom-to-top within the same diagonal.

This method works efficiently because it collects diagonals in a single pass and only iterates through all elements once more to produce the final output. It handles jagged arrays naturally, and there is no need for nested iteration over diagonals of varying lengths.

Approach Time Complexity Space Complexity Notes
Brute Force O(N + D^2) O(N + D) Collect elements by diagonal and reverse each; can be inefficient for jagged arrays with many diagonals
Optimal O(N) O(N) Hash map collects diagonals efficiently; single pass over all elements and then over diagonals to form result

Here, N is the total number of elements (sum(len(row) for row in nums)) and D is the maximum diagonal index.

Algorithm Walkthrough

  1. Initialize an empty dictionary diagonals to map diagonal index d = row + col to a list of elements.
  2. Initialize a variable max_diag to track the largest diagonal index seen.
  3. Iterate over each row r and each column c within that row.
  4. Compute the diagonal index d = r + c.
  5. Append nums[r][c] to diagonals[d].
  6. Update max_diag if d is larger than the current maximum.
  7. Initialize an empty list result to store the final diagonal traversal.
  8. Iterate from d = 0 to max_diag.
  9. For each diagonal list in diagonals[d], append its elements in reversed order to result.
  10. Return result.

Why it works: Each element is assigned to exactly one diagonal key based on row + col, which ensures proper grouping. Reversing each diagonal ensures the bottom-up order required by the problem. Iterating diagonals in increasing order guarantees the correct traversal sequence.

Python Solution

from typing import List
from collections import defaultdict

class Solution:
    def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]:
        diagonals = defaultdict(list)
        max_diag = 0
        
        for r, row in enumerate(nums):
            for c, val in enumerate(row):
                diagonals[r + c].append(val)
                max_diag = max(max_diag, r + c)
        
        result = []
        for d in range(max_diag + 1):
            result.extend(reversed(diagonals[d]))
        
        return result

This implementation first populates a hash map with elements grouped by their diagonal index. Each diagonal list is then reversed and appended to the final result. Using defaultdict simplifies the logic by avoiding explicit checks for missing keys.

Go Solution

func findDiagonalOrder(nums [][]int) []int {
    diagonals := make(map[int][]int)
    maxDiag := 0
    
    for r, row := range nums {
        for c, val := range row {
            diagonals[r+c] = append(diagonals[r+c], val)
            if r+c > maxDiag {
                maxDiag = r + c
            }
        }
    }
    
    result := []int{}
    for d := 0; d <= maxDiag; d++ {
        diag := diagonals[d]
        for i := len(diag) - 1; i >= 0; i-- {
            result = append(result, diag[i])
        }
    }
    
    return result
}

In Go, a map stores diagonal lists, and elements are appended similarly. Reversing each diagonal during iteration produces the required bottom-to-top order. Nil slices are handled naturally since appending to a nil slice creates a new slice.

Worked Examples

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

r c val d=r+c diagonals
0 0 1 0 {0: [1]}
0 1 2 1 {0:[1],1:[2]}
0 2 3 2 {0:[1],1:[2],2:[3]}
1 0 4 1 {0:[1],1:[2,4],2:[3]}
1 1 5 2 {0:[1],1:[2,4],2:[3,5]}
1 2 6 3 {0:[1],1:[2,4],2:[3,5],3:[6]}
2 0 7 2 {0:[1],1:[2,4],2:[3,5,7],3:[6]}
2 1 8 3 {0:[1],1:[2,4],2:[3,5,7],3:[6,8]}
2 2 9 4 {0:[1],1:[2,4],2:[3,5,7],3:[6,8],4:[9]}

Reversing each diagonal and concatenating: [1,4,2,7,5,3,8,6,9].

Complexity Analysis

Measure Complexity Explanation
Time O(N) Each element is processed once to populate diagonals and once when building the result
Space O(N) Diagonal hash map stores all elements; final result also stores all elements

The algorithm scales linearly with the total number of elements, making it efficient even for large jagged arrays.

Test Cases

# Provided examples
assert Solution().findDiagonalOrder([[1,2,3],[4,5,6],[7,8,9]]) == [1,4,2,7,5,3,8,6,9] # square matrix
assert Solution().findDiagonalOrder([[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]) == [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16] # jagged

# Edge cases
assert Solution().findDiagonalOrder([[1]]) == [1] # single element
assert Solution().findDiagonalOrder([[1],[2],[3]]) == [1,2,3] # single column
assert Solution().findDiagonalOrder([[1,2,3]]) == [1,2,3] # single row
assert Solution().findDiagonalOrder([[1,2],[3]]) == [1,3,2] # small jagged
Test Why
[[1,2,3],[4,5,6],[7,8,9]] Square matrix traversal correctness
[[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]] Handles jagged rows and larger input
[[1]] Minimal input
[[1],[2],[3]] Single column
[[1,2,3]] Single row
[[1,2],[3]] Small jagged array, tests diagonal ordering

Edge Cases

One important edge case is a single-row or single-column array. The algorithm handles this by naturally computing the diagonal indices; no reversal causes issues with single elements.

Another edge case is a highly jagged array,