498. Diagonal Traverse

• python3

Problems

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

Pseudocode

1. Set two pointers: one for the row and one for the column.
2. Create an array for the result.
3. Check the conditions. If the sum of the row and colum indices is even, move up in the matrix. If the sum is odd, move down.
4. Return the result array.

Code

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        
        r = 0
        c = 0
        m = len(mat)
        n = len(mat[0])
        flag = True
        result = [0] * (m*n)

        for i in range(0, m*n):
            result[i] = mat[r][c]
            if (r + c) % 2 == 0: # moving up
                if c == n - 1: 
                    r+=1
                elif r == 0 :   
                    c+=1
                else:
                    r-=1
                    c+=1
            else:                # moving down
                if r == m - 1:
                    c+=1
                elif c == 0:    
                    r+=1
                else: 
                    r+=1 
                    c-=1
        
        return result

Result

• Time Complexity: O(m*n) as we traverse the result array only once.
• Space Complexity: O(m*n) as we create an one array.