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.

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

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

Solution:

class Solution(object):
    def findDiagonalOrder(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: List[int]
        """
        res = []
        if not matrix or not matrix[0]:
            return res
        m = len(matrix)
        n = len(matrix[0])
        scan_nums = m + n - 1
        for i in range(0, scan_nums):
            if i % 2 == 0:
                x = i if i < m else m - 1
                y = 0 if i < m else i - m + 1
                while (x >= 0 and y < n):
                    res.append(matrix[x][y])
                    x -= 1
                    y += 1
            else:
                x = 0 if i < n else i - n + 1
                y = i if i < n else n - 1
                while (x < m and y >= 0):
                    res.append(matrix[x][y])
                    x += 1
                    y -= 1
        return res