Subarray Sum Equals K
整数の配列 nums と整数 k が与えられます。要素の合計がちょうど k になる部分配列の数を数えてください。部分配列とは、隣り合う1つ以上の要素が連続したものです。値が同じでも、開始位置または終了位置が異なる場合は、別々の部分配列として数えます。値は負またはゼロの場合もあります。
関数
- numsinteger-array
- 負の値やゼロを含むことがある整数の配列
- kinteger
- カウントされるために部分配列が到達しなければならない合計値
- 戻り値integer
- 要素の合計が k になる部分配列の数
制約
1 ≤ nums.length ≤ 2 × 104-1000 ≤ nums[i] ≤ 1000-107 ≤ k ≤ 107- この長さの配列の部分配列は最大でも200,010,000個なので、答えは32ビット符号付き整数に収まります。
例
- 入力
- nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
- 出力
- 4
- 説明
- 4つの連続区間の合計は7になります。
[3, 4]、[1, 3, 3]、[3, 3, 1]、そして[3, 4, -7, 1, 3, 3]です。最後の例では、-7が3と4を相殺し、その後合計が再び7まで増えるため、合計がkを超えた後でも連続区間が一致することがあります。
- 入力
- nums = [1, -1, 0]k = 0
- 出力
- 3
- 説明
- 合計が 0 になる部分配列は 3 つあります：
[1, -1]、[0]、そして配列全体の[1, -1, 0]です。[-1, 0]の並びの合計は -1 なので、数に含めません。
- 入力
- nums = [2, 2, 2]k = 4
- 出力
- 2
- 説明
- インデックス0と1にある連続部分列
[2, 2]と、インデックス1と2にある連続部分列[2, 2]は同じ値を持っていますが、位置が異なるため、どちらも数えます。配列全体の合計は6です。
提出時に隠しテスト+17件
発展問題
合計が k になる最長の部分配列の長さを、引き続き O(n) 時間で返すには、解法をどのように変更すればよいでしょうか？
ヒント
1つずつ開いてください。開くたびに少しずつ答えに近づきます。
すべての部分配列を調べる方法は機能しますが、20,000個の数値には約2億個の部分配列があります。値が負になることもあるため、スライディングウィンドウも機能しません。一度計算した数値を使って、任意の部分配列の合計を表せますか？
累積和を更新し続けます。2つの位置の間にある要素の合計は、終点での累積和から始点の直前の累積和を引いた値です。したがって、ここで終わる部分配列の合計がちょうど
kになるのは、以前の累積和が現在の累積和からkを引いた値と等しいときです。各プレフィックス和とその出現回数を対応付けたハッシュマップを使って、配列を一度だけ走査します。空のプレフィックスから始めます。つまり、和は 0 で、出現回数は 1 回です。各要素で、
prefix - kに格納されている回数を答えに加え、その後で現在のプレフィックスを記録します。
解説
n 個の数値からなる配列には n(n+1)/2 個の部分配列があり、n = 2 × 10^4 のとき約 2 × 10^8 個になるため、それぞれを足し合わせる方法では遅すぎます。また、負の値があるため、スライディングウィンドウも使えません。ウィンドウの合計は減ったあと再び増えることがあるので、いつ縮めるべきかを判断するルールがないからです。この問題を解く鍵は、すべての部分配列の合計を2つの累積和の差として表すことです。現在の要素で終わり、合計が k になる部分配列の数を数えるには、現在の累積和から k を引いた値と等しい過去の累積和の数を数えればよく、ハッシュマップを使えば1回の走査で答えられます。
それぞれ、合計を記録して始める
正しいが、最大のテストでは終わらない
考え方
すべての部分配列には、最初のインデックスstartと最後のインデックスendがあります。すべてのペアを調べてその合計を確認すれば、各部分配列をちょうど1回ずつ調べることになるので、数え方は正しいです。
各部分配列の合計を求めるために、3つ目のループは必要ありません。startを固定し、endを1つずつ右に進めながら、nums[end]を累積するtotalに加えます。totalには常にstartからendまでの要素の合計が入っているので、各部分配列の処理は加算1回と比較1回で済みます。
合計がkに達したり超えたりしても、処理を止めてはいけません。後の負の値によって、合計が再び小さくなることがあります。最初の例では、インデックス0からの合計は3、7、0、1、4、7と変化するため、開始位置が同じ部分配列で、インデックス5のときにも一致します。
処理コストはペアの数に比例します。n = 2 × 10^4の場合、ペアは約2 × 10^8個あり、Cなら問題ありませんが、Python、Ruby、Rでは遅すぎます。
アルゴリズム
countを0に設定する。- 0からn-1までの各
startについて、totalを0に設定する。 startからn-1までの各endについて、nums[end]をtotalに加える。totalがkと等しい場合、countに1を加え、いずれの場合も処理を続ける。countを返す。
def subarraySum(nums, k):
n = len(nums)
count = 0
for start in range(n):
total = 0
# Grow the subarray that begins at start, one element at a time
for end in range(start, n):
total += nums[end]
if total == k:
count += 1
return countカウントマップを使った累積和
考え方
prefix[j]を、最初のj個の要素の合計とします。空の接頭部分についてはprefix[0] = 0です。インデックスiからインデックスj-1までの部分配列の合計はprefix[j] - prefix[i]です。したがって、現在の要素で終わる部分配列の合計がちょうどkになるのは、それより前の接頭和が現在の接頭和からkを引いた値と等しい場合です。そのような接頭和が現れるたびに、一致する部分配列の開始位置が1つ示されます。
配列を1回だけ走査します。累積接頭和と、各接頭和が何回現れたかを記録するハッシュマップseenを保持します。各要素について、まずseen[prefix - k]を個数に加え、その後で現在の接頭和を記録します。記録する前に検索することで、部分配列が空になるのを防ぎます。k = 0の場合、先に記録すると現在の接頭和が自分自身と一致してしまいます。
k = 7の場合の最初の例を見てみましょう。接頭和は0、3、7、0、1、4、7、8、4です。インデックス1の後で接頭和が7になると、マップには0が1つあり、これによって[3, 4]が見つかります。インデックス5の後で再び7になると、マップには0が2つあります。空の接頭部分と、-7の後の接頭部分で、それぞれ[3, 4, -7, 1, 3, 3]と[1, 3, 3]が同時に見つかります。インデックス6の後で8になると、マップには1が1つあり、これによって[3, 3, 1]が見つかります。合計は4です。
マップに0が1回出現すると最初に設定しておくことで、インデックス0から始まる部分配列も数えられます。集合ではなく個数を記録するマップが重要なのは、同じ接頭和が繰り返し現れることがあり、そのたびに異なる部分配列が始まるためです。各要素につき検索と更新を1回ずつ行うので、時間計算量はO(n)であり、マップが保持するキーは最大でn+1個です。
アルゴリズム
seenというマップを作成し、seen[0] = 1とし、prefixとcountを0に設定します。- 各要素を
prefixに加算します。 - 存在しないキーは0として読み取り、
seen[prefix - k]をcountに加算します。 seen[prefix]に1を加算します。countを返します。
def subarraySum(nums, k):
# seen[p] = how many prefixes so far add up to p; the empty prefix adds up to 0
seen = {0: 1}
prefix = 0
count = 0
for num in nums:
prefix += num
# Each earlier prefix equal to prefix - k starts a subarray that ends here and sums to k
count += seen.get(prefix - k, 0)
seen[prefix] = seen.get(prefix, 0) + 1
return count
落とし穴と境界ケース
誤答の多くは、すべての値が正だとみなしたり、2つのマップ操作の順序を間違えたりすることから生じます。
- 合計が
kを超えた時点で縮めるスライディングウィンドウは、負の値があるとうまくいきません。最初の例では、4ではなく2を返します。インデックス6で合計が7を超えるまでウィンドウの左端がインデックス0のままなので、[1, 3, 3]や[3, 3, 1]を試しません。 seen[0] = 1を省くと、インデックス0から始まるすべての部分配列を見落とします。nums = [5]、k = 5の場合、1ではなく0を返します。- 検索する前に現在の累積和を記録すると、
kが0のときに空の部分配列を数えてしまいます。[1, -1, 0]の場合、3ではなく6を返します。 - 出現回数を記録するマップの代わりに累積和の集合を使うと、重複を過少に数えてしまいます。
[0, 0, 0]、k = 0の場合、同じ累積和のそれぞれの過去の出現が異なる部分配列の始まりになるため、答えは6です。 - 総和が
kを超えたときにブルートフォースの内側のループを抜けるのは、スライディングウィンドウの場合と同じ理由で誤りです。
よくある質問4
Subarray Sum Equals K の時間計算量は何ですか？
累積和とハッシュマップを使う解法は、O(n) 時間、O(n) の追加領域で実行できます。1 回の走査で、各要素につき 1 回の検索と 1 回の更新を行います。累積合計を使ってすべての部分配列を調べるには O(n²) 時間がかかり、各部分配列の合計を最初から計算すると O(n³) かかります。
なぜ Subarray Sum Equals K ではスライディングウィンドウが機能しないのでしょうか？
スライディングウィンドウは、ウィンドウを広げると合計が増え、狭めると減ることを前提としていますが、これはすべての値が正の場合にのみ成り立ちます。負の値があると、合計がすでに大きすぎるウィンドウでも、さらに広げることで条件に合う場合があるため、左端をいつ動かすべきかを判断するルールがありません。すべての値が正であれば、スライディングウィンドウを使って O(n) 時間、O(1) 空間で解けます。
ハッシュマップでは、なぜ 0 が 1 に対応する状態で始まるのでしょうか？
そのエントリは、最初の要素より前にある空の接頭辞を表し、その合計は 0 です。インデックス 0 から始まる部分配列の合計は、現在の接頭辞和からその空の接頭辞を引いた値になるため、このエントリがないと、それらの部分配列は一度もカウントされません。nums = [5]、k = 5の場合、5 - 5 = 0を検索するとそのエントリが見つかり、1が返されます。
部分配列の合計が K に等しいかどうかを、追加の O(1) 空間で解けますか？
1回の走査ではできません。ある要素で終わる一致数を数えるには、その要素より前にどの累積和があったかを知る必要があり、その種類は最大で n+1 個あります。マップがなければ、O(n²) の累積合計に頼ることになります。すべての値が正のときは、スライディングウィンドウを使うと、O(n) 時間、O(1) 空間で部分配列を数えられます。
似た問題
同じ考え方を使う問題です。2〜3問解くとパターンが身につきます。
Python
def subarraySum(nums, k):
# ここにコードを書いてくださいケース1
ケース2
ケース3
入力
nums = [3, 4, -7, 1, 3, 3, 1, -4] k = 7
期待値
4