C++で総和がKに等しい部分配列の個数を求める方法

整数型の配列 nums と整数 k が与えられたとき、要素の総和がちょうど k に等しくなる「連続する部分配列（サブアレイ）」の総数を求めることを考えます。たとえば nums = [1, 1, 1]、k = 2 の場合、条件を満たすのは [1, 1]（インデックス 0〜1）と [1, 1]（インデックス 1〜2）の 2 つなので、答えは 2 となります。

解法の考え方：累積和とハッシュマップ

すべての部分配列を総当たりで調べる方法もありますが、O(n²) の計算量がかかり非効率です。そこで役立つのが「累積和（プレフィックスサム）」のテクニックです。ある位置で終わる部分配列の和が k になるのは、その位置までの累積和から、より手前の位置までの累積和を引いた差が k に等しいときです。つまり「現在の累積和 − k」が過去に何回現れたかを数えれば、それがそのまま答えになります。この出現回数をハッシュマップで管理することで、配列を一度走査するだけで解けます。

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

• マップ sums を定義し、temp := 0、sums[0] := 1、ans := 0 として初期化する
• i を 0 から配列サイズ − 1 までループする
  • temp := temp + n[i] として累積和を更新する
  • sums にキー (k − temp) が存在する場合は、ans := ans + sums[k − temp] を行う
  • sums[−temp] の値を 1 増やす
• 最後に ans を返す

初期化で sums[0] = 1 としているのは、配列の先頭から始まる部分配列も正しくカウントできるようにするためです。また、この実装ではマップのキーに累積和の符号を反転した値（−temp）を登録しているため、探索も k − temp という形になっています。よく見られる「sums[temp − k] を参照する」標準的な書き方と、結果は完全に同一です。

例（C++コード）

理解を深めるために、実際の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int subarraySum(vector<int>& n, int k) {
      unordered_map <int, int> sums;
      int temp = 0;
      sums[0] = 1;
      int ans = 0;
      for(int i = 0; i < n.size(); i++){
         temp += n[i];
         if(sums.find(k - temp) != sums.end()){
            ans += sums[k - temp];
         }
         sums[-temp]++;
      }
      return ans;
  }
};
main(){
   Solution ob;
   vector<int> v = {1,1,1};
   cout << (ob.subarraySum(v, 2));
}

入力

[1,1,1]
2

出力

2

計算量

このアルゴリズムは配列を一度だけ走査するため、時間計算量は O(n) です。累積和の出現回数を最大で n 種類保持するため、空間計算量も O(n) となります。二重ループですべての部分配列の和を確認する素朴な O(n²) のアプローチと比べて、大きな入力でも高速に動作します。