Diagonal Traverse

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:

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        m = len(mat)
        n = len(mat[0])
        op = []
        x, y = 0, 0
        while len(op) != m * n:
            slope = (x + y) & 1
            op.append(mat[y][x])
            dx, dy = (-1, 1) if slope else (1, -1)
            if x + dx < 0 or x + dx >= n or y + dy < 0 or y + dy >= m:
                if slope:
                    if y < m - 1:
                        y += 1
                    else:
                        x += 1
                else:
                    if x < n - 1:
                        x += 1
                    else:
                        y += 1
            else:
                x += dx
                y += dy
        return op