对于每条对角线,行号 i 加列号 j 是一个定值。示例 1 正中间对角线的 i+j 恒为 2。

设 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] 加入答案。

#
# 代码中的类名、方法名、参数名已经指定,请勿修改,直接返回方法规定的值即可
#
# 
# @param mat int整型二维数组 
# @return int整型一维数组
#
class Solution:
    def diagonalOrder(self , mat: List[List[int]]) -> List[int]:
        # write code here
        m, n = len(mat), len(mat[0])
        ans = []
        for k in range(m+n-1):
            if k%2:
                i = 0 if k<n else k-n+1
                j = k if k<n else n-1
                while i<m and j>=0:
                    ans.append(mat[i][j])
                    i += 1
                    j -= 1
            else:
                i = k if k<m else m-1
                j = 0 if k<m else k-m+1
                while i>=0 and j<n:
                    ans.append(mat[i][j])
                    i -= 1
                    j += 1
        return ans

复杂度分析

  • 时间复杂度:O(mn),其中 m 和 n 分别为 grid 的行数和列数。
  • 空间复杂度:O(1)。返回值不计入。