498. Diagonal Traverse

Given a matrix of M x N elements (M rows, N columns), return all elements of the matrix in diagonal order as shown in the below image.

Example:

Input:

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

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

Explanation:

Thoughts:

Check boundary, if OOB then reset the proper initial point + switching the traverse direction

class Solution(object):
    def findDiagonalOrder(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: List[int]
        """
        if len(matrix) == 0 or len(matrix[0]) == 0: return []

        m, n = len(matrix), len(matrix[0])

        row, col, d = 0 , 0 , 0

        dirs = [-1,1, 1,-1] # upright vs downleft
        res = []
        for i in range(m*n):
            res.append(matrix[row][col])
            row+= dirs[2*d + 0]
            col+= dirs[2*d + 1]
            # check if OOB, then switching the direction + initial point
            if row >= m: row, col, d = m - 1, col + 2, 1 - d
            if col >= n: row, col, d = row + 2, n - 1, 1 - d
            if row < 0: row, col, d = 0, col, 1 - d
            if col < 0: row, col, d = row, 0, 1 - d

        return res