C++で合計がSとなるバイナリ部分配列を数える方法

0と1のみから構成される配列Aが与えられたとき、合計がSとなる空でない（非空の）部分配列がいくつ存在するかを求める問題です。例えば、入力が [1,0,1,0,1]、S = 2 の場合、答えは 4 になります。該当する部分配列は [1,0,1]、[1,0,1,0]、[0,1,0,1]、[1,0,1] の4つです。

解法のアプローチ

この問題は「スライディングウィンドウ（尺取り法）」を用いた補助関数 atMost() を使うことで効率的に解けます。atMost(A, x) は「合計が x 以下となる部分配列の個数」を返す関数です。すると、合計がちょうど S となる部分配列の個数は、次の差として求められます。

atMost(A, S) − atMost(A, S − 1)

アルゴリズムの手順

• 配列Aと整数xを受け取るメソッド atMost() を定義します。

• x < 0 の場合は 0 を返します。j := 0、ret := 0 と初期化します。

• i を 0 から配列Aのサイズまでループします。

  • x から A[i] を減算します。

  • x < 0 の間、x に A[j] を加算し、j を 1 増やします。

  • ret に i − j + 1 を加算します。

• ret を返します。

• メイン処理では以下を実行します。

• ret := atMost(A, S) − atMost(A, S − 1)

• ret を返します。

この手法では、各要素が最大でも2回しか走査されないため、時間計算量は O(n)、追加の空間計算量は O(1) という非常に効率的な実装になります。

C++での実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    int atMost(vector<int>& A, int x){
        if(x < 0) return 0;
        int j = 0;
        int ret = 0;
        for(int i = 0; i < A.size(); i++){
            x -= A[i];
            while(x < 0){
                x += A[j];
                j++;
            }
            ret += i - j + 1;
        }
        return ret;
    }
    int numSubarraysWithSum(vector<int>& A, int S) {
        return atMost(A, S) - atMost(A, S - 1);
    }
};
main(){
    vector<int> v1 = {1,0,1,0,1};
    Solution ob;
    cout << (ob.numSubarraysWithSum(v1, 2));
}

入力

[1,0,1,0,1]

出力

4