Pythonでk×k部分行列の最小値を求めるプログラム【スライディングウィンドウ法】

問題の概要

2次元の行列と整数 k が与えられたとき、すべての k × k 部分行列に含まれる最小値を要素とする新しい行列を返すプログラムを考えます。

たとえば、次のような入力が与えられたとします。

ここで k = 2 とした場合、出力は [[3, 5], [3, 3]] になります。

この結果は、各2×2の部分行列ごとの最小値を表しています。具体的には次の通りです。

• 左上の部分行列：最小値は 3
  

  3 5
  8 6

• 右上の部分行列：最小値は 5
  

  5 6
  6 5

• 左下の部分行列：最小値は 3
  

  8 6
  4 3

• 右下の部分行列：最小値は 3
  

  6 5
  3 12

解き方のアプローチ

この問題は、スライディングウィンドウ最小値（Sliding Window Minimum）のテクニックを使うことで効率的に解けます。両端キュー（deque）を利用すれば、各行・各列について線形時間でウィンドウ内の最小値を求められます。

手順は以下の通りです。

1. 行方向の処理：各行に対して、サイズ k のスライディングウィンドウで最小値を求め、その結果で行を置き換えます。
2. 列方向の処理：各列に対しても同様に、サイズ k のスライディングウィンドウで最小値を求め、その結果で列を置き換えます。
3. 結果の抽出：処理後の行列から、位置 (i + k − 1, j + k − 1) の値を取り出して、答えとなる行列 ret を構築します。

アルゴリズムの詳細

• 各行 r について：
  • 空の両端キュー q と空のリスト nrow を用意します。
  • i を 0 から行の長さまで順に処理します。
    • q の先頭のインデックスが i − k と等しい場合（ウィンドウから外れた要素）、先頭を削除します。
    • q の末尾のインデックスに対応する値が row[i] より大きい間、末尾を削除します。
    • i を q の末尾に追加し、現在のウィンドウの最小値 row[q[0]] を nrow に追加します。
  • matrix[r] を nrow で置き換えます。
• 各列 j についても同じ処理を行い、得られた ncol の値で列を更新します。
• (行数 − k + 1) × (列数 − k + 1) のサイズのリスト ret を作成し、ret[i][j] = matrix[i + k − 1][j + k − 1] を代入します。
• ret を返します。

実装例

理解を深めるために、実際のPythonコードを見てみましょう。

import collections
class Solution:
    def solve(self, matrix, k):
        for r, row in enumerate(matrix):
            q = collections.deque()
            nrow = []
            for i in range(len(row)):
                if q and q[0] == i - k:
                    q.popleft()
                while q and row[q[-1]] > row[i]:
                    q.pop()
                q.append(i)
                nrow.append(row[q[0]])
            matrix[r] = nrow
        for j in range(len(matrix[0])):
            q = collections.deque()
            ncol = []
            for i in range(len(matrix)):
                if q and q[0] == i - k:
                    q.popleft()
                while q and matrix[q[-1]][j] > matrix[i][j]:
                    q.pop()
                q.append(i)
                ncol.append(matrix[q[0]][j])
            for i in range(len(matrix)):
                matrix[i][j] = ncol[i]
        ret = [[0] * (len(matrix[0]) - k + 1) for _ in range(len(matrix) - k + 1)]
        for i in range(len(ret)):
            for j in range(len(ret[0])):
                ret[i][j] = matrix[i + k - 1][j + k - 1]
        return ret
ob = Solution()
print(ob.solve(matrix = [
    [3, 5, 6],
    [8, 6, 5],
    [4, 3, 12]
], k = 2))

入力

[[3, 5, 6],[8, 6, 5],[4, 3, 12]], 2

出力

[[3, 5], [3, 3]]

計算量について

このアルゴリズムの時間計算量は O(R × C) です（R は行数、C は列数）。各要素は最大でも一度キューに追加され、一度削除されるだけなので、全体として線形時間で処理できます。一方、各部分行列ごとに k² 個の要素を総当たりで調べる素朴な方法では O(R × C × k²) かかるため、この手法を用いることで大幅な高速化が可能になります。