C++で合計がKで割り切れる部分配列の個数を効率的に求める方法

整数型の配列 A が与えられたとき、要素の合計が K で割り切れる「空でない連続する部分配列」の個数を求めます。例えば、A = [4,5,0,-2,-3,1]、k = 5 の場合、出力は 7 になります。該当する7つの部分配列は [[4,5,0,-2,-3,1], [5], [5,0], [5,0,-2,-3], [0], [0,-2,-3], [-2,-3]] です。

解法のアプローチ

この問題は「累積和」と「剰余」を組み合わせたハッシュマップを使うことで、O(n) の計算量で効率的に解くことができます。考え方の基本は次の通りです。2つの累積和が K で割った余りを等しく持つならば、その間に対応する部分配列の合計は必ず K の倍数になります。そこで、各剰余がこれまで何回現れたかを記録し、同じ剰余が出るたびに組み合わせの数を加算していきます。

• ハッシュマップ m を作成し、m[0] を 1 に初期化します（先頭から始まる部分配列も正しくカウントするためです）
• temp := 0、ans := 0、n := 配列 a のサイズとします
• i を 0 から n − 1 まで繰り返します
  • temp := temp + a[i]（累積和を更新）
  • x := (temp mod k + k) mod k（負の数にも対応できる剰余計算）
  • ans := ans + m[x]（同じ剰余を持つ過去の位置との組み合わせをカウント）
  • m[x] を 1 増やします
• ans を返します

ここでのポイントは (temp % k + k) % k という式です。C++では被除数が負の場合、剰余演算の結果が負になることがあるため、この式を使うことで常に非負の剰余を安全に取得できます。

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

実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int subarraysDivByK(vector<int>& a, int k) {
      unordered_map <int, int> m;
      m[0] = 1;
      int temp = 0;
      int ans = 0;
      int n = a.size();
      for(int i = 0; i < n; i++){
         temp += a[i];
         int x = (temp % k + k) % k;
         ans += m[x];
         m[x]++;
      }
      return ans;
   }
};
main(){
   vector<int> v = {4,5,0,-2,-3,1};
   Solution ob;
   cout <<(ob.subarraysDivByK(v, 5));
}

入力

[4,5,0,-2,-3,1]
5

出力

7

この手法により、全ての部分配列を総当たりで調べる O(n²) 以上のアプローチと比べ、線形時間 O(n) で答えを得られるため、大きな入力サイズでも高速に動作します。