Pythonで最大K回のインクリメント操作後に等しい要素からなる最長部分リストを求めるプログラム

問題の概要

数値のリスト nums と整数 k が与えられます。「リスト内の任意の1つの要素を1だけ増やす」という操作を最大 k 回まで行えるとき、すべての要素が等しい値になるような最長の部分リスト(連続する部分列)の長さを求めます。

たとえば、入力が nums = [3, 5, 9, 6, 10, 7]、k = 6 の場合を考えてみましょう。9 を1回、6 を4回インクリメントすれば、部分リスト [10, 10, 10] が作れるため、答えは 3 になります。

解法のステップ

この問題は、スライディングウィンドウと単調デック(モノトニックデック)を組み合わせることで効率的に解けます。手順は以下のとおりです。

• nums が空の場合は 0 を返します。

• wMax を nums と同じサイズの両端キュー(deque)として初期化し、ペア (nums[0], 0) を挿入します。

• i := 0、inc := 0 とします。

• j を 1 から nums のサイズまで1ずつ増やしながら、以下を繰り返します。

  • wMax が空でなく、wMax[0][1] < i である間、wMax の左端の要素を削除します。

  • pMax := wMax[0][0] とします。

  • wMax が空でなく、末尾要素の値が nums[j] 以下である間、右端の要素を削除します。

  • ペア (nums[j], j) を wMax の末尾に追加します。

  • pMax < wMax[0][0] の場合は、inc += (j - i) * (wMax[0][0] - pMax) とします。

  • それ以外の場合は、inc += pMax - nums[j] とします。

  • inc > k になった場合は、次のようにウィンドウを縮めます。

    • inc -= wMax[0][0] - nums[i] とします。

    • wMax が空でなく、wMax[0][1] <= i である間、左端の要素を削除します。

    • wMax[0][0] < nums[i] の場合は、inc -= (nums[i] - wMax[0][0]) * (j - i) とします。

    • i := i + 1 とします。

• 最後に len(nums) - i を返します。

実装例

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

サンプルコード

from collections import deque
class Solution:
   def solve(self, nums, k):
      if not nums:
         return 0
      wMax = deque([(nums[0], 0)], maxlen=len(nums))
      i = 0
      inc = 0
      for j in range(1, len(nums)):
         while wMax and wMax[0][1] < i:
            wMax.popleft()
         pMax = wMax[0][0]

         while wMax and wMax[-1][0] <= nums[j]:
            wMax.pop()
         wMax.append((nums[j], j))
         if pMax < wMax[0][0]:
            inc += (j - i) * (wMax[0][0] - pMax)
         else:
            inc += pMax - nums[j]
         if inc > k:
            inc -= wMax[0][0] - nums[i]
            while wMax and wMax[0][1] <= i:
               wMax.popleft()
            if wMax[0][0] < nums[i]:
               inc -= (nums[i] - wMax[0][0]) * (j - i)
            i += 1
      return len(nums) - i
ob = Solution()
nums = [3, 5, 9, 6, 10, 7]
k = 6
print(ob.solve(nums, k))

入力

[3, 5, 9, 6, 10, 7], 6

出力

3

アルゴリズムのポイント

この実装では、変数 inc が「現在のウィンドウ内の全要素をウィンドウの最大値に揃えるために必要なインクリメントの合計回数」を表しています。デック wMax は各要素とそのインデックスのペアを保持しており、ウィンドウ内の最大値を常に先頭から O(1) で参照できるように管理されています。

inc が k を超えた時点で、ウィンドウの左端 i を進めてコストを再計算し、条件を満たす範囲を維持します。最終的な答えは「リストの長さから左端の位置を引いた値」、つまり有効なウィンドウの最大幅となります。全体の時間計算量は O(n)、空間計算量も O(n) であり、非常に効率的な解法です。