https://leetcode.com/problems/diagonal-traverse

Main trick to solve this problem easily is to understand the structure of our diagonal traverse. Let us call level each diagonal we traverse:

  1. In first level we have only one coordinates: [0, 0].
  2. In second level we have two points: [0, 1] and [1, 0].
  3. In third level we have three points: [0, 2], [1, 1] and [2, 0].

Now, aogorithm becomes very easy:

  1. For each level put all elements in this level to dictionary levels: note, that for each level we put it in direct order.
  2. Now, for each level choose if we traverse it in direct order or reverse, using [::lev%2*2-1] notation. If lev is odd number, then we have [::1], that is direct order, if lev is even, we have [::-1], that is reverse order.

Complexity: time complexity is O(mn): we traverse each element once. Additional space complexity in this approach is also O(mn), because we keep levels dictionary. Theoretically, we can put our data directly into long list and then reverse some parts, we will have O(1) additional space in this case.

class Solution:
    def findDiagonalOrder(self, matrix):
        if not matrix: return []
        m, n = len(matrix), len(matrix[0])
        levels = defaultdict(list)
        for i, j in product(range(m), range(n)):
            levels[i+j].append(matrix[i][j])
                
        out = []
        for lev in range(m + n):
            out += levels[lev][::lev%2*2-1]   
        return out

If you like the solution, you can upvote it on leetcode discussion section: Problem 0498