Skip to content

498. Diagonal Traverse

Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

Example 1:

img

Input: mat = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,2,4,7,5,3,6,8,9]

Example 2:

Input: mat = [[1,2],[3,4]]
Output: [1,2,3,4]

Constraints:

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

Solution:

对于每条对角线，行号\(i\) 加列号 \(j\) 是一个定值。示例 1 正中间对角线的 \(i + j\) 恒为 2。（可以回想一下 51. N 皇后 的写法）

设 \(k = i + j\)，那么左上角那条对角线的 $ k=0$ ，右下角那条对角线的 \(k=(m−1)+(n−1)=m+n−2\)。

枚举 \(k=0,1,2,…,m+n−2\)，就相当于在从左上到右下，一条一条地枚举对角线。

由于 \(i+j=k\)，知道 \(j\) 就知道 \(i\)，所以我们只需要计算出每条对角线的\(j\) 的最小值和最大值，就可以开始遍历对角线了。

由于 \(j=k−i\)，当 \(i=m−1\) 时\(j\) 取到最小值 \(k−m+1\)，但这个数不能是负数，所以最小的\(j\) 是 \(max(k−m+1,0)\)。 由于 \(j=k−i\)，当 \(i=0\) 时 \(j\) 取到最大值 \(k\)，但这个数不能超过 \(n−1\)，所以最大的\(j\) 是 \(min(k,n−1)\)。 然后就可以模拟了：

枚举 \(k=0,1,2,…,m+n−2\)。 如果 \(k\) 是偶数，我们从小到大枚举 \(j\)，否则从大到小枚举 \(j\)。其中 \(j\) 的范围是 \([max(k−m+1,0),min(k,n−1)]\)。 把$ mat[k−j][j]$ 加入答案。

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

        int[] result = new int[m * n];
        int index = 0;
        // (0,0)  k = 0
        // (1,1)  k = 2
        // (2,2)  k = 3
        // i, j
        // i + j = k
        int i = 0;
        for (int k = 0; k <= (n - 1) + (m - 1); k++){
            int minJ = Math.max(k - (n - 1), 0); // i = n - 1, j = k - (n -1)
            int maxJ = Math.min(k, m - 1); // i = 0, j = k

            if (k % 2 == 0){
                // 偶数从左到右
                for (int j = minJ; j <= maxJ; j++){
                    result[i] = mat[k - j][j];
                    i++;
                }
            }else {
                for (int j = maxJ; j >= minJ; j--){
                    result[i] = mat[k - j][j];
                    i++;
                }
            }
        }

        return result;

    }
}

// TC: O(mn)
// SC: O(1)