C++で0・1・2の個数が等しい部分文字列を数える方法

0、1、2のみで構成された文字列 str が与えられたとき、0・1・2がそれぞれ同数含まれる部分文字列をすべて見つけ、その総数を求めるのが本記事の目的です。例えば str が「12012」の場合、条件を満たす部分文字列は「120」「201」「012」の3つであるため、答えは3となります。

具体例で理解する

入力： str="112200120"

出力： 0・1・2の個数が等しい部分文字列の数： 5

説明： 条件を満たす部分文字列は以下の5つです。

str[0-5]="112200"、str[1-6]="122001"、str[5-7]="012"、str[6-8]="120"、str[0-8]="112200120"

入力： str="12012"

出力： 0・1・2の個数が等しい部分文字列の数： 3

説明： 条件を満たす部分文字列は以下の3つです。

str[0-2]="120"、str[1-3]="201"、str[2-4]="012"

プログラムで使うアプローチ

この問題を効率的に解く鍵となるのは「位置ごとの個数の差分」です。位置 i までに読み込んだ 0・1・2 の個数をそれぞれ c0、c1、c2 とすると、部分文字列 str[j..i] が条件を満たすのは、「位置 j−1 における差分ペア (c0−c1, c0−c2)」と「位置 i における差分ペア」が一致するとき、そしてそのときに限られます。したがって、各位置で差分ペアを計算し、同じペアが過去に何回出現したかを map で管理しながら足し合わせていけば、答えを求められます。

• 対象となる文字列を受け取り、その長さを求めます。
• データを関数に渡して処理を行います。
• 条件を満たす部分文字列の個数を格納するための一時変数 count を用意します。
• 差分ペアとその出現頻度を対応付ける map 型の変数を作成します。
• ペア (0, 0) に 1 を設定し（空の接頭辞に相当）、FORループを 0 から文字列の長さまで回します。
• ループ内では、str[i] が '0' なら 0 のカウントを、'1' なら 1 のカウントを、それ以外なら 2 のカウントをそれぞれインクリメントします。
• c0−c1 および c0−c2 を計算し、そのペアをキーとして map を参照します。
• count に map_pair[差分ペア] の現在の値を加算し、その後 map_pair の該当エントリを 1 増やします。
• ループ終了後、count を返します。
• 結果を出力します。

実装例

#include <bits/stdc++.h>
using namespace std;

// 0・1・2の個数が等しい部分文字列の数を数える関数
int countSubstrings(string str){
    int n = str.length();
    int count = 0;
    int c0 = 0, c1 = 0, c2 = 0;
    // 差分ペアと出現回数を対応付けるマップ
    map<pair<int,int>, int> mp;
    mp[{0, 0}] = 1; // 空の接頭辞に相当する初期値
    for(int i = 0; i < n; i++){
        if(str[i] == '0'){
            c0++;
        }
        else if(str[i] == '1'){
            c1++;
        }
        else{
            c2++;
        }
        pair<int,int> p = {c0 - c1, c0 - c2};
        count += mp[p]; // 同じ差分ペアの過去の出現回数を加算
        mp[p]++;
    }
    return count;
}
int main(){
    string str = "112200120";
    cout<<"Count of Substrings with equal number of 0s, 1s and 2s are: "<<countSubstrings(str);
    return 0;
}

出力

上記のコードを実行すると、次の出力が得られます。

Count of Substrings with equal number of 0s, 1s and 2s are: 5

計算量について

各位置で map への参照と更新を1回ずつ行うため、全体の時間計算量は O(n log n)、空間計算量は O(n) となります。std::unordered_map とカスタムハッシュ関数を組み合わせれば、平均 O(n) まで高速化できます。すべての部分文字列を素朴に走査する方法では O(n²) 以上の時間がかかることを考えると、非常に大きな改善といえます。