Diagonal Traverse of 2D Array

Problem Description

Given a 2D integer array nums, return all elements of nums in diagonal order as shown in the below images.

Example 1:

Input: nums = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,4,2,7,5,3,8,6,9]

Example 2:

Input: nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
Output: [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

Naive Approach

A naive approach would be to iterate through the diagonals. The challenge here is to determine the start and end points of each diagonal. We can iterate through the matrix and store elements in a dictionary (or hash map) where the key is the diagonal index (row + col). Then, we iterate through the diagonal indices and collect the elements in the correct order.

Code (Python)

def diagonal_traverse_naive(nums):
    diagonal_map = {}
    for i in range(len(nums)):
        for j in range(len(nums[i])):
            diagonal_index = i + j
            if diagonal_index not in diagonal_map:
                diagonal_map[diagonal_index] = []
            diagonal_map[diagonal_index].append(nums[i][j])

    result = []
    diagonal_keys = sorted(diagonal_map.keys())
    for key in diagonal_keys:
        result.extend(diagonal_map[key][::-1]) # Reverse order for each diagonal

    return result

Time Complexity

O(M * N), where M is the number of rows and N is the number of columns. We iterate through each element of the matrix once.

Space Complexity

O(M * N) because, in the worst case, we might store all elements of the matrix in the diagonal_map.

Optimal Approach

The optimal approach focuses on collecting the elements directly in the desired order without sorting. We can iterate through diagonals identified by the sum row + col, and then iterate backwards (from larger row index to smaller row index) along each diagonal to collect elements.

Code (Python)

def diagonal_traverse(nums):
    result = []
    n = len(nums)
    
    for diagonal_index in range(sum(len(row) for row in nums) - 1):
        for i in range(n):
            j = diagonal_index - i
            if 0 <= i < n and 0 <= j < len(nums[i]):
                result.append(nums[i][j])
    
    return result

Another possible solution using collections.defaultdict:

from collections import defaultdict

def diagonal_traverse_defaultdict(nums):
    diagonals = defaultdict(list)
    for i in range(len(nums)):
        for j in range(len(nums[i])):
            diagonals[i + j].append(nums[i][j])
    
    result = []
    for key in sorted(diagonals.keys()):
        result.extend(diagonals[key][::-1])
    
    return result

Time Complexity

O(M * N), where M is the number of rows and N is the number of columns. We still iterate through all elements.

Space Complexity

O(1) excluding the space for the output array, if we are very careful, as we are not using additional data structures that scale with the input size significantly (in the first optimal approach). The defaultdict solution is O(M*N).

Edge Cases

• Empty Input: If the input nums is empty, return an empty list.
• Uneven Rows: The matrix may have rows of different lengths, which must be handled correctly.
• Single Element: If the matrix contains only one element, return a list containing that element.