【Python】合計がターゲット値と一致するサブリストの個数を効率的に求める方法

数値のリスト nums とターゲット値 target が与えられたとき、要素の合計が target と一致するサブリスト(連続する部分列)がいくつ存在するかを求める問題について解説します。

たとえば、nums = [3, 0, 3]、target = 3 という入力の場合、答えは 4 になります。これは、合計が 3 になるサブリストとして [3]、[3, 0]、[0, 3]、[3] の 4 つが存在するためです。

解法のアプローチ:累積和とハッシュマップ

すべてのサブリストを総当たりで調べると O(n²) の時間がかかりますが、累積和(プレフィックスサム)と辞書(ハッシュマップ)を組み合わせることで、O(n) で効率的に解くことができます。

考え方はシンプルです。先頭から現在位置までの累積和を s とすると、「現在の累積和 − target」が過去のいずれかの時点の累積和と一致していれば、その間の区間の合計は target と一致することになります。そこで、各累積和の出現回数を辞書に記録しておき、一致するたびにその回数を答えに加算します。

アルゴリズムの手順

• temp:累積和の出現回数を記録する空のマップ(辞書)を用意する
• temp[0] := 1:初期状態(累積和 0)を 1 回として登録する
• s := 0(累積和)、ans := 0(答え)で初期化する
• i を 0 から nums のサイズ未満まで繰り返す:
  • s := s + nums[i] で累積和を更新する
  • comp := s - target を計算する
  • comp が temp に存在すれば、ans := ans + temp[comp] を実行する
  • temp[s] := temp[s] + 1 で現在の累積和を記録する
• ans を返す

Python実装コード

from collections import defaultdict

class Solution:
    def solve(self, nums, target):
        temp = defaultdict(int)
        temp[0] = 1
        s = 0
        ans = 0
        for i in range(len(nums)):
            s += nums[i]
            comp = s - target
            if comp in temp:
                ans += temp[comp]
            temp[s] += 1
        return ans

ob = Solution()
nums = [3, 0, 3]
target = 3
print(ob.solve(nums, target))

入力

[3, 0, 3], 3

出力

4

まとめ

この手法では、リストを一度走査するだけで答えが得られるため、時間計算量は O(n) に抑えられます。また、0 や負の数が含まれるリストでも正しく動作します。defaultdict(int) を使うことで、キーが存在しない場合の初期化処理も簡潔に書ける点がポイントです。「Subarray Sum Equals K」など、競技プログラミングやコーディング面接で頻出の定番テクニックなので、ぜひ覚えておきましょう。