Algorithm Problem 003: Diagonal Traverse II

Posted on In Word count in article: 769 Reading time ≈ 3 mins.

Problem

Problem Source： 1424. Diagonal Traverse II

Problem Description：

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

Example 1:

image-20231029115351760

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

Example 2:

image-20231029115127660

1
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]

Example 3:

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

Constraints:

  • 1 <= nums.length <= 105
  • 1 <= nums[i].length <= 105
  • 1 <= sum(nums[i].length) <= 105
  • 1 <= nums[i][j] <= 105

Approach

For a natural number $k \in \mathbb{N}$, the coordinates of the elements on the $(k + 1)$-th diagonal are given by $\set{(i,j) | i + j = k}$. For example, when $k = 2$, the coordinates of the elements on the 3rd diagonal are $\set{(0,2),(1,1),(2,0)}$.

If nums were a matrix, this problem would be a simple simulation problem, just traversing it diagonally.

However, given that the number of rows $m \le 10^5$ and columns $n \le 10^5$, a time complexity of $O(mn)$ would be too slow. Therefore, we should not treat nums as a matrix for traversal to avoid accessing many empty cells.

Let’s observe example 1:

image-20231029115351760

Observation1️⃣: For any $k \in \mathbb{N}$, the elements on the $(k + 1)$-th diagonal are traversed from bottom-left to top-right, i.e., $(k, 0), (k - 1,1),\cdots,(0,k)$. Therefore, if we want to obtain the elements on the $(k + 1)$-th diagonal, we should traverse each element in the corresponding rows of nums in reverse order.

Observation2️⃣: It is evident that for each row nums[j] in nums, the elements in this row belong to different diagonals, making them independent of each other.

Observation3️⃣: When traversing in reverse order:

  • For the last row with elements $12,13,14,15,16$, they are the first elements on their respective diagonals, even though the diagonals that $13,14,15,16$ belong to do not have elements in the first column.
  • For the second-to-last row with elements $9,10,11$, they are the next elements to be traversed on their respective diagonals.
  • This pattern continues…

Our goal is to collect all elements from nums diagonally. To achieve this, we can first create a list for each diagonal and add the corresponding elements to that list. Then, we traverse these lists in order to obtain the final result. If we follow this approach, we need to create a total of m + n - 1 lists.

1
2
3
4
5
6
7
8
9
10
11
int m = nums.size();
// n is the maximum number of columns
int n = nums.stream().mapToInt(List::size).max().getAsInt();
// ans stores m + n - 1 lists, representing elements traversed on different diagonals
List<List<Integer>> ans = new ArrayList<>();
IntStream.range(0, m + n - 1).forEach(i -> ans.add(new ArrayList<>()));

// Key problem: How to traverse??

// Traverse nums in forward order for each row
return ans.stream().flatMap(l -> l.stream()).mapToInt(i -> i).toArray();

So, how do we traverse? To avoid timeout, we need to traverse row by row.

  • According to ‘Observation 1️⃣’, traverse nums in reverse order.
  • According to ‘Observation 2️⃣’, for each row, for each traversed element nums[i][j], put it into the $(i + j + 1)$-th list.
  • According to ‘Observation 3️⃣’, append the element nums[i][j] to the corresponding list. This is because the previous $n - i - 1$ rows ($n -1, n - 2 , \cdots, i + 1$) have already been traversed for the diagonal, and their elements are already stored in the list in the same order —— we can prove this using recursion or mathematical induction.

The code is straightforward:

1
2
3
4
5
6
7
8
// Traverse nums in reverse
for (int i = m - 1; i >= 0; i--) {
    List<Integer> row = nums.get(i);
    // For each element nums[i][j], put it into the (i + j + 1)-th list
    for (int j = 0; j < row.size(); j++) {
        ans.get(i + j).add(row.get(j));
    }
}

The time complexity is $O(m + n + \text{(# elements in nums)})$.

The space complexity is $O(m + n)$.

Code

Java

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int m = nums.size();
        // n is the maximum number of columns
        int n = nums.stream().mapToInt(List::size).max().getAsInt();
        // ans stores m + n - 1 lists, representing elements traversed on different diagonals
        List<List<Integer>> ans = new ArrayList<>();
        IntStream.range(0, m + n - 1).forEach(i -> ans.add(new ArrayList<>()));

        // Traverse nums in reverse
        for (int i = m - 1; i >= 0; i--) {
            List<Integer> row = nums.get(i);
            // For each element nums[i][j], put it into the (i + j + 1)-th list
            for (int j = 0; j < row.size(); j++) {
                ans.get(i + j).add(row.get(j));
            }
        }
        // Traverse nums in forward order for each row
        return ans.stream().flatMap(l -> l.stream()).mapToInt(i -> i).toArray();
    }
}