Diagonal Traverse

Simulation Array

Problem

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.

For example:

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

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

Thought Process

• Notice how essentially, we are only going in two directions, up and down. More specifically, going in a up-right pattern and a down-left pattern

• When going up, if you hit the top border, you have to go right one element; if you hit the far right border, you have to go down one element. Similarly, when going down, if you hit the bottom border, you have to go right one element; if you hit the far left border, you have to go down one element.

• The reason a boolean works better for keeping track of direction compared to a variable (like in Spiral Order Matrix) is because we only need to keep track of two directions compared to four. The boolean allows us to flip between traversing the two directions like a light switch.

Solution

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        #only going in two directions, up and down
        #use a boolean expression to flip between the directions
        
        if len(matrix) == 0:
            return
        
      
        bottom = len(matrix)-1
        left = 0
        right = len(matrix[0])-1
        goDown = False
        res = []
        row, col = 0,0
        
        while row <= bottom and col <= right:
            res.append(matrix[row][col])
            if goDown:
                if row == bottom or col == 0:
                    goDown = False
                    if row == bottom:
                        col+=1
                    else:     
                        row+=1
                else:
                    row+=1
                    col-=1
            else:
                if row == 0 or col == right:
                    goDown = True
                    if col == right:
                        row+=1
                    else:
                        col+=1
                else:
                    row-=1
                    col+=1
        return res
                
                
                    
           

Key Facts

• Two cases we need to check for each direction:

  • When going down if we are in the first column or the last row

  • When going up if we are in the last column or top row

• To keep track of changing directions, we use a boolean variable since there are only two directions we are traversing. All we need to check is if this boolean is true or false to determine the direction we are traversing

Time Complexity

• Time: $O(n)$

• Space: $O(n)$for the return array