Problem description:

Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

Example 1:

https://assets.leetcode.com/uploads/2021/04/10/diag1-grid.jpg

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

Example 2:

1
2
Input: mat = [[1,2],[3,4]]
Output: [1,2,3,4]

Constraints:

  • m == mat.length
  • n == mat[i].length
  • 1 <= m, n <= 104
  • 1 <= m * n <= 104
  • 105 <= mat[i][j] <= 105

Solution:

The key here is to realize that the sum of indices on all diagonals are equal.For example, in

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

2, 4 are on the same diagonal, and they share the index sum of 1. (2 is matrix[0][1] and 4 is in matrix[1][0]). 3,5,7 are on the same diagonal, and they share the sum of 2. (3 is matrix[0][2], 5 is matrix[1][1], and 7 is matrix [2][0]).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        level = defaultdict(list)
        for i in range(len(mat)):
            for j in range(len(mat[0])):
                level[i+j].append(mat[i][j])
        
        res = []
        for i, entry in enumerate(level.values()):
            if i % 2 == 0:
                res.extend(entry[::-1])
            else:
                res.extend(entry)
        return res

time complexity: $O(mn)$
space complexity: $O(mn)$
reference:
related problem: