Pythonで最大k個の0を反転して最長の連続する1の長さを求めるプログラム

0と1のみで構成されるバイナリリストと、整数kが与えられたとします。最大k個の0を1に変更できるとき、すべてが1で構成される最も長い連続区間(部分リスト)の長さを求めるのがこの問題です。

例えば、入力が nums = [0, 1, 1, 0, 0, 1, 1]、k = 2 の場合、出力は 6 になります。中央にある2つの0を1に変更すれば、リストは [0, 1, 1, 1, 1, 1, 1] となり、連続する6つの1が得られるためです。

解決アプローチ:スライディングウィンドウ法

この問題は「スライディングウィンドウ(尺取り法)」と呼ばれる手法を使うことで、O(n)の計算量で効率的に解くことができます。考え方はシンプルで、右端を伸ばしながらウィンドウ内の0の数がkを超えないよう左端を調整し、その都度ウィンドウ幅の最大値を記録していきます。

具体的な手順は以下の通りです。

• 変数を初期化します:zeros = 0(現在のウィンドウ内の0の数)、ans = 0(答え)、j = 0(ウィンドウの左端インデックス)
• リストの各要素について、インデックスiと値nを使って以下を繰り返します。
  • nが0であれば、zerosを1増やします。
  • zerosがkより大きい間、次の処理を行います。
    • nums[j]が0であれば、zerosを1減らします。
    • jを1増やしてウィンドウの左端を右へ移動させます。
  • 現在のウィンドウ幅 i - j + 1 がansより大きければ、ansを更新します。
• 最後にansを返します。

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

実装例

class Solution:
    def solve(self, nums, k):
        zeros = 0
        ans = 0
        j = 0
        for i, n in enumerate(nums):
            zeros += n == 0
            while zeros > k:
                zeros -= nums[j] == 0
                j += 1
            if i - j + 1 > ans:
                ans = i - j + 1
    return ans

ob = Solution()
nums = [0, 1, 1, 0, 0, 1, 1]
k = 2
print(ob.solve(nums, k))

入力

[0, 1, 1, 0, 0, 1, 1], 2

出力

6

このアルゴリズムでは、各要素が高々2回(右端として1回、左端として1回)走査されるだけなので、全体の計算量はO(n)、追加のメモリ使用量はO(1)で済みます。バイナリ配列に関する「最大k回の操作で〜を最大化する」という種類の問題には、このスライディングウィンドウのパターンが広く応用できるので、ぜひ覚えておきましょう。