Leetcode Diagonal Traverse problem solution

In this Leetcode Diagonal Traverse problem solution Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

Problem solution in Python.

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        if not matrix or not matrix[0]:
           return []
        m,n = len(matrix),len(matrix[0])
        result = []
        for i in range(m+n-1):
            tmp = []
            if i < n:
                row,col = 0,i
            else:
                row,col = i-n+1,n-1
            while row < m and col >=0:
                tmp.append(matrix[row][col])
                row += 1
                col -= 1
            if i%2:
                result.extend(tmp)
            else:
                result.extend(tmp[::-1])
        return result

Problem solution in Java.

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix.length == 0) return new int[]{};
        int m = matrix.length, n = matrix[0].length;
        int a = 0, b = 0, i = 0, step = 1;
        int[] res = new int[m*n];
        
        while (i < res.length) {
            while (a >= 0 && a < m && b >= 0 && b < n) {
                res[i++] = matrix[a][b];
                a -= step;
                b += step;
            }
            a += step; b -= step;
            if ((a == 0 || a == m-1) && b < n-1) {
                b++;
            } else if (a < m && (b == 0 || b == n-1)) {
                a++;
            }
            step = -step;
        }
        return res;
    }
}

Problem solution in C++.

vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        int m = matrix.size(), n;
        if(m!=0)n = matrix[0].size();
        vector<int> vec;
        int i=0, j=0, dir = 1; //dir=1 means move right-up, dir = 0 means move bottom-left
        while(i<m && j<n && i>=0 && j>=0){
            vec.push_back(matrix[i][j]);
            if(i==m-1&&j==n-1)break;
            if(dir){
                if(i-1>=0&&j+1<n){i--,j++;}  //move right-up
                else if(j+1<n){j++,dir = 0;}  // move right
                else {i++,dir = 0;}  //move down
            }
            else{
                if(i+1<m && j-1>=0){i++,j--;}  //move bottom-left
                else if(i+1<m){i++, dir = 1;}  //move down
                else {j++,dir = 1;} //move right
            }
        }
        return vec;
    }