Pythonでバイナリリスト内の合計がkとなるサブリストの個数を求めるプログラム

問題の概要

0と1のみで構成されるバイナリリストが与えられたとします。さらに別の入力として整数 k が与えられ、要素の合計がちょうど k に一致するサブリスト(連続する部分配列)の個数を求める必要があります。

例えば、入力が nums = [1, 0, 0, 1, 1, 1, 0, 1]、k = 3 の場合、出力は 8 になります。これは、条件を満たすサブリストとして [1,0,0,1,1]、[0,0,1,1,1]、[0,0,1,1,1,0]、[0,1,1,1]、[0,1,1,1,0]、[1,1,1]、[1,1,1,0]、[1,1,0,1] の8つが存在するためです。

解決のためのアプローチ

この問題は「累積和」と「ハッシュマップ(辞書)」を組み合わせることで効率的に解けます。リストを走査しながら各位置までの累積和を記録し、「現在の累積和 − k」という値がこれまでに何回出現したかを数えることで、合計が k になるサブリストの数を求めます。手順は以下の通りです。

• sums := キー 0 に値 1 を持つマップ(辞書)で初期化する
• r_sum := 0(累積和を保持する変数)
• ans := 0(答えを保持する変数)
• nums の各要素 x に対して以下を繰り返す
  • r_sum := r_sum + x(累積和を更新)
  • ans := ans +(r_sum − k が sums に存在すればその値、なければ 0)
  • sums[r_sum] := 現在の値 + 1(累積和の出現回数をカウントアップ)
• 最後に ans を返す

この手法により、すべてのサブリストを総当たりする O(n²) の方法に比べて、線形時間 O(n) で処理が完了します。

実装例

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

def solve(nums, k):
   sums = {0: 1}
   r_sum = 0
   ans = 0

   for x in nums:
      r_sum += x
      ans += sums.get(r_sum - k, 0)
      sums[r_sum] = sums.get(r_sum, 0) + 1

   return ans

nums = [1, 0, 0, 1, 1, 1, 0, 1]
k = 3
print(solve(nums, k))

入力

[1, 0, 0, 1, 1, 1, 0, 1], 3

出力

8

計算量について

このアルゴリズムの時間計算量は O(n)、空間計算量は O(n) です(n はリストの長さ)。累積和の出現回数を辞書で管理することで、各要素を一度だけ走査すればよいため、大規模なデータでも高速に動作します。