Neunomizuの日記

俺だけが俺だけじゃない

# LeetCode Medium 498. Diagonal Traverse

tags: leetcode

問題

イデア

M x Nの行列が与えられたとき,対角線方向にジグザグに動いた時に通る要素を返せ」という問題です.

このような問題で重要なのは問題を小分けすることで,ここでは最初に右上から左下への対角線方向に進んだ際に通る要素を考え,次にジグザグに動くことを考えます.

解法

  • { [0, 0] } -> { [0, 1], [1, 0] } -> { [0, 2], [1, 1], [2, 0] } -> ...という風に進むと考えます
    • この時,ある規則で出現する配列群([0, 1], [1, 0]などの塊)の行と列の添字の和は1ずつ増えていき,最大の和は行列の行と列の和-1です.
    • また,もう一点注目するべきことがあります.3 x 3行列で考えます.
      • 上の例の後に,{ [1, 2], [2, 1] } -> { [2, 2] }となります.
      • つまり,行の添字が行列の行の長さよりも大きくなると数字のパターンが変わります.
  • これを考えれば,後は配列群の行と列の添字の和で答えの配列に逆さに加えるかどうかを考えれば良いです.

計算量

  • 時間計算量
    • $O(NM)$
  • 空間計算量
    • $O(NM)$

Python

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        if not matrix or not matrix[0]:
            return []
        height, width = len(matrix), len(matrix[0])
        result = []
        
        for diag in range(height + width - 1):
            tmp = []
            row, col = 0, 0
            if diag < width:
                row, col = 0, diag
            else:
                row, col = diag - width + 1, width - 1
            
            while row < height and col > -1:
                tmp.append(matrix[row][col])
                row += 1
                col -= 1
            
            if diag % 2 == 0:
                result.extend(tmp[::-1])
            else:
                result.extend(tmp)
        return result