498. Diagonal Traverse

class Solution {
    private static final int[][] DIRECTIONS = new int[][]{{-1, 1}, {1, -1}};

    public int[] findDiagonalOrder(int[][] mat) {
        int m = mat.length, n = mat[0].length;

        int[] res = new int[m * n];
        int index = 0;
        int directionIndex = 0;

        // starting point: left to right
        for (int startColIndex = 0; startColIndex < n; startColIndex++) {
            int startRowIndex = 0;

            int startIndex = index, endIndex = index;
            int i = startRowIndex, j = startColIndex;
            while (i >= 0 && i < m && j >= 0 && j < n) {
                endIndex = index;
                res[index++] = mat[i++][j--];
            }

            if (directionIndex == 0) {
                reverse(res, startIndex, endIndex);
            }
            directionIndex = directionIndex == 0 ? 1 : 0;
        }

        // starting point: top to bottom
        for (int startRowIndex = 1; startRowIndex < m; startRowIndex++) {
            int startColIndex = n - 1;

            int startIndex = index, endIndex = index;
            int i = startRowIndex, j = startColIndex;
            while (i >= 0 && i < m && j >= 0 && j < n) {
                endIndex = index;
                res[index++] = mat[i++][j--];
            }

            if (directionIndex == 0) {
                reverse(res, startIndex, endIndex);
            }
            directionIndex = directionIndex == 0 ? 1 : 0;
        }

        return res;
    }

    private void reverse(int[] nums, int left, int right) {
        while (left < right) {
            swap(nums, left++, right--);
        }
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

以 i = 0 这条边及 j = n - 1 这条边为起点，统一向左下方向遍历并添加，并对需要逆向的遍历进行 reverse 即可。

Reference

498. Diagonal Traverse