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:

Note:

The total number of elements of the given matrix will not exceed 10,000.

## Solution

from typing import List

class Solution:
def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
if not matrix:
return []

result = []

# Row and Col Directions
directions = [
[-1, 1],
[1, -1],
]

# Index to directions
curDir = 0
curRow = 0
curCol = 0
numRows = len(matrix)
numCols = len(matrix[0])

for i in range(numRows * numCols):
result.append(matrix[curRow][curCol])

# Top right corner
if curRow + directions[curDir][0] < 0 and curCol + directions[curDir][1] == numCols:
curRow += 1
curDir = (curDir + 1) % 2
# Bottom left corner
elif curCol + directions[curDir][1] < 0 and curRow + directions[curDir][0] == numRows:
curCol += 1
curDir = (curDir + 1) % 2
# Left or right border
elif curCol + directions[curDir][1] == numCols or curCol + directions[curDir][1] < 0:
curRow += 1
curDir = (curDir + 1) % 2
# Top or bottom border
elif curRow + directions[curDir][0] == numRows or curRow + directions[curDir][0] < 0:
curCol += 1
curDir = (curDir + 1) % 2
else:
curRow += directions[curDir][0]
curCol += directions[curDir][1]

return result


## Test Cases

test = Solution()
[1,2,3],
[4,5,6],
[7,8,9],
])
[1,2,3],
[4,5,6],
])
[1,2,3],
])
[1,2],
[4,5],
[7,8],
])