Pythonでサイズkのサブリストの最大値を求めるアルゴリズム

問題の概要

リスト nums と整数 k が与えられたとき、連続する k 個の要素からなる各サブリスト（スライディングウィンドウ）ごとの最大値を求め、その結果をリストとして返すことを考えます。

例えば、nums = [12, 7, 3, 9, 10, 9]、k = 3 の場合、出力は [12, 9, 10, 10] になります。
これは、先頭から順に3要素ずつ区切った [12, 7, 3]、[7, 3, 9]、[3, 9, 10]、[9, 10, 9] のそれぞれの最大値を並べたものです。

解法の考え方

この問題は、現在の最大値とその位置を記録しながら窓をずらしていくことで解けます。具体的な手順は以下の通りです。

• k が nums のサイズより大きい場合は、空のリストを返します。

• 結果を格納するためのリスト res を用意します。

• 最初の窓（先頭の k 個）について、最大値 temp とそのインデックス point を求めます。i を 0 から k − 1 まで走査し、nums[i] > temp であれば temp := nums[i]、point := i と更新します。

• temp を res の末尾に追加します。

• 続いて i を k から nums のサイズまで走査し、新しい要素が入ってきた際の挙動を次のように場合分けします。

  • nums[i] < temp かつ (i − point) < k の場合：
    現在の最大値はまだ窓の中にあるため、temp はそのまま維持されます（temp := nums[point]）。

  • nums[i] < temp かつ (i − point) ≥ k の場合：
    現在の最大値が窓の外に出てしまったため、point を i − k + 1 に設定し直し、j を i − k + 1 から i まで走査して窓内の新しい最大値を探します。見つかった位置の値を temp := nums[point] とします。

  • それ以外の場合：
    新しく入ってきた要素が現在の最大値以上であるため、temp := nums[i]、point := i と更新します。

• 各ステップで temp を res の末尾に追加していきます。

• 最後に res を返します。

実装例

それでは、実際のコードを見ながら理解を深めましょう。

class Solution:
    def solve(self, nums, k):
        if k > len(nums):
            return []
        res = []
        temp = nums[0]
        point = 0
        for i in range(k):
            if nums[i] > temp:
                temp = nums[i]
                point = i
        res.append(temp)
        for i in range(k, len(nums)):
            if nums[i] < temp and (i - point) < k:
                temp = nums[point]
            elif nums[i] < temp and (i - point) >= k:
                point = i - k + 1
                for j in range(i - k + 1, i + 1):
                    if nums[j] > nums[point]:
                        point = j
                temp = nums[point]
            else:
                temp = nums[i]
                point = i
            res.append(temp)
        return res

ob = Solution()
nums = [12, 7, 3, 9, 10, 9]
k = 3
print(ob.solve(nums, k))

入力

[12, 7, 3, 9, 10, 9], 3

出力

[12, 9, 10, 10]

計算量のポイント

このアルゴリズムは、最大値がまだ窓内に残っている間は再計算を行わないため、多くのケースで O(n) に近い効率で動作します。ただし、最大値が窓から外れるたびに窓内の再走査が発生するため、単調に増減するデータなどの最悪ケースでは O(n × k) になる可能性があります。より安定した性能が必要な場面では、両端キュー（collections.deque）を活用したスライディングウィンドウ最大値の手法も検討してみるとよいでしょう。