498. Diagonal Traverse

This is a hard problem.

If I got this problem the first time when I was interviewed, I would not be able to solve it.

    public int[] findDiagonalOrder(int[][] mat) {
        int m = mat.length, n = mat[0].length;
        int[] res = new int[m*n];
        int i=0, j=0;
        for(int k=0;k<m*n;k++){
            res[k] = mat[i][j];
            if((i+j)%2==0){  //when the arrow is going up
               if(j==n-1)    //must judge j==n-1 first
                    i++;
                 else if(i==0)
                    j++;
                else{        //when not on edge
                    i--;
                    j++;
                }
            }else{           //when the arrow is going down
                if(i==m-1)   //must judge i==m-1 first
                    j++;
                else if(j==0)
                    i++;
                else{        //when not on edge
                    i++;
                    j--;
                }
            }
        }
        return res;
    }