Pythonで最大値と最小値の差がk以下となる最長の部分リストの長さを求めるプログラム

数値のリスト nums と整数 k が与えられたとき、「リスト内の最大要素と最小要素の絶対差が k 以下」という条件を満たす、最も長い連続した部分リストの長さを求めることを考えます。

たとえば、入力が nums = [2, 4, 6, 10]、k = 4 の場合、出力は 3 になります。[2, 4, 6] を選べば、最大値 6 と最小値 2 の差がちょうど 4 となり、条件を満たすからです。

解き方のアプローチ

この問題は、スライディングウィンドウと単調な両端キュー(deque)を組み合わせることで効率的に解けます。ウィンドウの右端を拡張しながら、ウィンドウ内の最大値と最小値をそれぞれ専用の deque で管理し、両者の差が制限を超えたら左端を縮めていくのがポイントです。

アルゴリズムの手順

• 2つの両端キュー maxd(最大値管理用)と mind(最小値管理用)を作成する
• i := 0(ウィンドウの左端)、res := 1(答えの初期値)とする
• リスト A の各インデックス j と値 a について、以下を繰り返す:
  • maxd が空でなく、a が maxd の末尾より大きい間、末尾の要素を削除する
  • mind が空でなく、a が mind の末尾より小さい間、末尾の要素を削除する
  • a を maxd と mind の末尾に追加する
  • maxd[0] - mind[0] > limit である間、以下を繰り返す:
    • maxd[0] == A[i] なら、maxd の先頭を削除する
    • mind[0] == A[i] なら、mind の先頭を削除する
    • i を 1 増やして、ウィンドウの左端を進める
  • res を res と j - i + 1 の大きい方で更新する
• 最後に res を返す

Pythonでの実装例

from collections import deque
class Solution:
    def solve(self, A, limit):
        maxd = deque()
        mind = deque()
        i = 0
        res = 1
        for j, a in enumerate(A):
            while maxd and a > maxd[-1]:
                maxd.pop()
            while mind and a < mind[-1]:
                mind.pop()
            maxd.append(a)
            mind.append(a)
            while maxd[0] - mind[0] > limit:
                if maxd[0] == A[i]:
                    maxd.popleft()
                if mind[0] == A[i]:
                    mind.popleft()
                i += 1
            res = max(res, j - i + 1)
        return res
ob = Solution()
nums = [2, 4, 6, 10]
k = 4
print(ob.solve(nums, k))

入力

[2, 4, 6, 10], 4

出力

3

計算量について

各要素は deque に対して高々 1 回の追加と 1 回の削除しか行われないため、全体の時間計算量は O(n)、空間計算量も O(n) となります。すべての部分リストを素朴に全探索する O(n²) のアプローチと比べて、大幅に高速に動作するのが特徴です。