498. Diagonal Traverse

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

/**
 * r+c determines which diagonal you are on. For ex: [2,0],[1,1],[0,2] are all
 * on same diagonal with r+c =2. If you check the directions of diagonals, first
 * diagonal is up, second diagonal is down third one is up and so on.. therefore
 * (r+c)%2 simply determines direction.
 *
 * Time Complexity: O(M*N)
 *
 * Space Complexity: O(1) without considering result space.
 *
 * M = Number of rows. N = Number of columns.
 */
class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
            return new int[0];
        }

        int r = 0;
        int c = 0;
        int rows = matrix.length;
        int cols = matrix[0].length;

        int[] result = new int[rows * cols];

        for (int i = 0; i < result.length; i++) {
            result[i] = matrix[r][c];

            if ((r + c) % 2 == 0) { // Move Up
                if (c == cols - 1) {
                    // Reached last column. Now move to below cell in the same column.
                    r++;
                } else if (r == 0) {
                    // Reached first row. Now move to next cell in the same row.
                    c++;
                } else {
                    // Somewhere in middle. Keep going up diagonally.
                    r--;
                    c++;
                }
            } else { // Move Down
                if (r == rows - 1) {
                    // Reached last row. Now move to next cell in same row.
                    c++;
                } else if (c == 0) {
                    // Reached first columns. Now move to below cell in the same column.
                    r++;
                } else {
                    // Somewhere in middle. Keep going down diagonally.
                    r++;
                    c--;
                }
            }
        }

        return result;
    }
}