Problem Statement in English

Given a matrix of integers mat, return all elements of the matrix in diagonal order as shown in the provided images.

Approach

This is a hardcore implementation problem that requires careful consideration of the matrix boundaries and the direction of traversal.

Solution in Python


class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        rows = len(mat)
        cols = len(mat[0])

        i,j = 0,0
        res = []

        up = True

        while 0 <= i < rows and 0 <= j < cols:
            res.append(mat[i][j])

            if up:
                if not i and j < cols - 1:
                    j += 1
                    up = False
                elif j == cols - 1:
                    i += 1
                    up = False
                else:
                    i -= 1
                    j += 1
            else:
                if not j and i < rows - 1:
                    i += 1
                    up = True
                elif i == rows - 1:
                    j += 1
                    up = True
                else:
                    i += 1
                    j -= 1
            
        return res

Complexity

  • Time: $O(m * n)$
    Since we traverse each element of the matrix exactly once.

  • Space: $O(1)$
    We are using only a constant amount of extra space.

And we are done.