Comment on page

# 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]
Note: The total number of elements of the given matrix will not exceed 10,000.
The Idea: First grab the elements that make the 'L' as shown below. Then take the diagonal of each element. Reverse the diagonal as you go for each alternative diagonal. Complexity: O(n) time and O(1) space
class Solution:
def findDiagonalOrder(self, matrix):
"""
:type matrix: List[List[int]]
:rtype: List[int]
"""
if not matrix:
return []
n_rows = len(matrix)
n_cols = len(matrix)
L = []
for i in range(0, n_rows):
L.append((i, 0))
for j in range(1, n_cols):
L.append((n_rows - 1, j))
def diagonal_along(r, c, reverse):
diagonal = []
while r >= 0 and c < len(matrix):
diagonal.append(matrix[r][c])
r -= 1
c += 1
if reverse:
return list(reversed(diagonal))
return diagonal
sol = []
reverse = False
for x, y in L:
sol += diagonal_along(x, y, reverse)
reverse = not reverse
return sol