LeetCode 498. 对角线遍历 java实现

/**
 * 给定一个含有 M x N 个元素的矩阵（M 行，N 列），
 * 请以对角线遍历的顺序返回这个矩阵中的所有元素，对角线遍历如下图所示。
 *
 */

public class FindDiagonalOrder498 {

    /**
     * 题解：
     *      1. 按照从右上到左下的方向取值，每个对角线上位置的行列之和是相同的，范围是0~m+n-1;
     *      2. 遍历这个和 sum，每次都从右上到左下；
     *          此处，可以将sum值的增加理解为：先遍历数组第一行，到最后一个数字之后，继续往下遍历。
     *      3. row + column = sum; 每次只需要判定row是不是超过数组的m：
     *          如果没有超过，则row都是从0开始增加1，此时column递减1；
     *          否则，row的值为 sum - m + 1;
     *      4. 由下往上，坐标为（i，0) （从最左一列开始）或者 (m-1, i-m + 1)(此时从上往下，均为从最下一行开始）
     *
     * @param mat
     * @return
     */
    public int[] findDiagonalOrder(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;
        int[] res = new int[m * n];
        int pos = 0;
        for (int i = 0; i < m + n - 1; i++) {
            if (i % 2 == 1) {
                int x = i < n ? 0 : i - n + 1;
                int y = i < n ? i : n - 1;
                while (x < m && y >= 0) {
                    res[pos] = mat[x][y];
                    pos++;
                    x++;
                    y--;
                }
            } else {
                int x = i < m ? i : m - 1;
                int y = i < m ? 0 : i - m + 1;
                while (x >= 0 && y < n) {
                    res[pos] = mat[x][y];
                    pos++;
                    x--;
                    y++;
                }
            }
        }
        return res;
    }
}