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[0])

        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[0]):
                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