Hex.*

LC498 - Diagonal Traverse

Posted on In
Leetcode: 129/2658 (53+62+14)

Description

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

Example 1:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,2,4,7,5,3,6,8,9]
Example 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

Assuming n = row, m = column, the current point is (i,j), since each traversal is on the line i+j=z, so write (i,z-i), z is from 0 to m+n-2

Use i to traverse, the point of each traversal is (i, z-i), there are two inequalities
0<=i<=n-1
0<=z-i<=m-1, that is, z-m+1<=i<=z

So the lower bound is: max(0,z-m+1)
The upper bound is: min(n-1,z)

Then traverse from the upper bound to the lower bound, and from the next to the upper bound, and both the upper and lower bounds can be obtained.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# O(nm) time | O(n) space
class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        row = len(mat)
        col = len(mat[0])
        # 0: bottom to top, 1: top to bottom
        direction = 0 
        res = []

        for line in range(row+col-1):
            if direction:
                for item in range(max(0, line - col +1), min(row-1, line)+1):
                    res.append(mat[item][line-item])
                direction = 0
            else:
                for item in range(min(row-1, line), max(0, line-col+1)-1, -1):
                    res.append(mat[item][line-item])
                direction = 1
        return res