Diagonal Traverse — LC Medium

Companies: Facebook, Microsoft

498. Diagonal Traverse

Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

Example 1:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,2,4,7,5,3,6,8,9]

Example 2:

Input: mat = [[1,2],[3,4]]
Output: [1,2,3,4]

Constraints:

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 104
• 1 <= m * n <= 104
• -105 <= mat[i][j] <= 105

Code:

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        
        # whenever we have 2d matrix, it is a good idea to think about the row and column indexes.
        maxRow = len(mat)
        maxCol = len(mat[0])
        myDict = collections.defaultdict(list)
        for row in range(maxRow):
            for col in range(maxCol):
                summed = row + col
                myDict[summed].append(mat[row][col])
                
        for key, lists in myDict.items():
            if key % 2 == 0:
                myDict[key] = myDict[key][::-1]
                
        return [i for i in myDict.values() for i in i]

Time complexity: O(N²)

Space complexity: O(N)