498. Diagonal Traverse

/**

 * 498. Diagonal Traverse

 * 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.

 *

 * Example:

 *

 * Input:

 * [

 *  [ 1, 2, 3 ],

 *  [ 4, 5, 6 ],

 *  [ 7, 8, 9 ]

 * ]

 *

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

 */

See the solution on github here: https://github.com/zcoderz/leetcode/blob/main/src/main/java/arrays/multidimensional/DiagonalTraverse.java

Observation: Questions is asking to print the diagonals in zigzag order. We can write the diagonals in a list and alternatively reverse the list to get zig zag order. To get the diagonals we can traverse from top right to bottom left. We can keep the diagonal index as an incrementing from 0 to sum of rows and columns. The starting column index is minimum of diagonal index and number of columns -1. Starting row index is simply diagonal index – column index. We loop via incrementing row and decrementing column. Reversal of items in diagonals needs to occur every other diagonal, we can run the reversal when diagonal % 2 ==0.

Here is the code:

public int[] findDiagonalOrder(int[][] matrix) {
    int rows = matrix.length;
    if (rows > 0) {
        int cols = matrix[0].length;
        int size = rows * cols;
        int[] out = new int[size];
        if (cols == 0) return out;
        int index = 0;

        List<Integer> vals = new ArrayList<>();

        for (int di = 0; di < rows + cols; di++) {
            int row = 0;
            int col;
            vals.clear();

            col = Math.min(di, cols - 1);
            if (col == cols - 1) {
                row = di - col;
            }

            while (row < rows && col >= 0) {
                vals.add(matrix[row][col]);
                row++;
                col--;
            }
            if (di % 2 == 0) {
                Collections.reverse(vals);
            }

            for (int val : vals) {
                out[index++] = val;
            }
        }
        return out;
    }
    return new int[0];
}