0498. Diagonal Traverse

https://leetcode.com/problems/diagonal-traverse

Description

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

ac

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        // edge cases
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0) return new int[0];
        int row = matrix.length;
        int col = matrix[0].length;
        int[] res = new int[row * col];

        int r = 0, c = 0;
        int index = 0;
        boolean up = true;
        while (r < row && c < col) {
            res[index++] = matrix[r][c];

            if (up) { // upward
                while (withinBoundary(r-1, c+1, matrix)) {
                    res[index++] = matrix[--r][++c];
                }
                if (withinBoundary(r, c+1, matrix)) c++;
                else if (withinBoundary(r+1, c, matrix)) r++;
                else break;
                up = !up;
            } else {  // downward
                while (withinBoundary(r+1, c-1, matrix)) {
                    res[index++] = matrix[++r][--c];
                }
                if (withinBoundary(r+1, c, matrix)) r++;
                else if (withinBoundary(r, c+1, matrix)) c++;
                else break;
                up = !up;
            }
        }
        // res[index] = matrix[row-1][col-1];  // add last one

        return res;
    }

    private boolean withinBoundary(int r, int c, int[][] matrix) {
        if (r < 0 || r >= matrix.length || c < 0 || c >= matrix[0].length)
            return false;
        return true;
    }
}