498. 对角线遍历

问题描述

给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。

示例 1：

输入：mat = [[1,2,3],[4,5,6],[7,8,9]]
输出：[1,2,4,7,5,3,6,8,9]

示例 2：

输入：mat = [[1,2],[3,4]]
输出：[1,2,3,4]

提示：

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 104
• 1 <= m * n <= 104
• -105 <= mat[i][j] <= 105

解题思路与代码实现

class Solution {

        /**
         * 解题思路：
         * 每条对角线的元素的行列下标之和是相等的：
         * 其为奇数时，向左下遍历；
         * 其为偶数时，向右上遍历
         * 需要注意在碰到行列边界时，需要相应右移或者下移
         */
        public int[] findDiagonalOrder(int[][] mat) {
            int m = mat.length, n = mat[0].length;
            int[] res = new int[m * n];
            int row = 0, col = 0;
            // 一共要访问m*n个元素
            for (int i = 0; i < res.length; i++) {
                res[i] = mat[row][col];
                // 偶数向右上方继续寻找，碰到边界需要调整
                if ((row + col) % 2 == 0) {
                    if (col == n - 1) {  // 下移
                        row++;
                    } else if (row == 0) { // 右移
                        col++;
                    } else {    // 右上
                        row--;
                        col++;
                    }
                } else {
                    if (row == m - 1) { // 右移
                        col++;
                    } else if (col == 0) {  // 下移
                        row++;
                    } else {    // 左下
                        row++;
                        col--;
                    }
                }
            }
            return res;
        }

    }