👨‍💻 738LeetCode

Diagonal Traverse

Question

Given a matrix of M x N elements, return all elements of the matrix in diagonal order.

Example

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

Solution

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        if not matrix:
            return []
        m, n = len(matrix), len(matrix[0])
        ret = [0] * (m * n)
        row, col, go_up = 0, 0, True
        for i in range(m * n):
            ret[i] = matrix[row][col]
            row += -1 if go_up else 1
            col += 1 if go_up else -1
            if row >= m:
                row, col, go_up = m - 1, col + 2, not go_up
            if col >= n:
                row, col, go_up = row + 2, n - 1, not go_up
            if row < 0:
                row, go_up = 0, not go_up
            if col < 0:
                col, go_up = 0, not go_up
        return ret