C++でバイナリ配列内の0のみ・1のみからなる部分配列を数える方法

0と1のみで構成される配列 arr[] が与えられたとき、各部分配列が「0のみ」または「1のみ」を含むような部分配列の総数を求めるのが本記事の目的です。

例えば、配列が [1,0,0] の場合、部分配列は以下のようになります。

• 0のみを含む部分配列:[0]、[0]、[0,0]
• 1のみを含む部分配列:[1]

具体例で理解しよう

入力: arr[] = { 0, 0, 1, 1, 1, 0 }

出力:

0のみを含む部分配列の数: 4
1のみを含む部分配列の数: 6

解説: 各部分配列は以下の通りです。

0のみの場合: [0], [0], [0], [0,0] の計4つ ( arr[0], arr[1], arr[5], arr[0-1] )
1のみの場合: [1], [1], [1], [1,1], [1,1], [1,1,1] の計6つ ( arr[2], arr[3], arr[4],
arr[2-3], arr[3-4], arr[2-4] )

入力: arr[] = { 1, 0, 1, 0 }

出力:

0のみを含む部分配列の数: 2
1のみを含む部分配列の数: 2

解説:

0のみの場合: [0], [0] の計2つ ( arr[1], arr[3] )
1のみの場合: [1], [1] の計2つ ( arr[0], arr[2] )

アルゴリズムの考え方

このプログラムでは、0のみを含む部分配列と1のみを含む部分配列をそれぞれ独立に数えるため、配列を2回走査します。連続する0の個数と連続する1の個数を記録するカウンター count_0 および count_1 を用意します。

ここで重要なポイントは、「長さ n の同一要素の連続列」から作れる部分配列の総数は n×(n+1)/2 という組み合わせの公式で求められることです。例えば、連続する1が3個ある場合、[1]、[1]、[1]、[1,1]、[1,1]、[1,1,1] の6つ = 3×4/2 となります。

この値を配列の末尾に達するまで合計していき、total_0(および total_1)として累積します。

手順の詳細

• 数値の配列 arr[] を用意します。
• 関数 sub_zero_one(int arr[], int size) は、配列を受け取り、0のみを含む部分配列の数と1のみを含む部分配列の数を返します。
• 部分配列の累積カウントとして、初期値 temp_0 と temp_1 を用意します。
• 連続する0と1の一時的な個数として count_0 と count_1 を用意します。
• forループを使い、i=0 から i<size まで配列を走査します。
• 現在の要素が0であれば count_0 をインクリメントします。
• そうでなければ、それまでの count_0 個の0から作れる部分配列の数 temp_one_0 = count_0 × (count_0 + 1) / 2 を計算します。
• この値をこれまでの temp_0 に加算し、count_0 をリセットします。
• 同様の手順を、変数 count_1、temp_one_1、temp_1 を使って1についても実行します。
• 両方の走査が終わった時点で、temp_0 と temp_1 にはそれぞれ0のみ・1のみを含む部分配列の総数が格納されています。

C++実装例

#include <bits/stdc++.h>
using namespace std;
void sub_zero_one(int arr[], int size){
    int count_1 = 0;
    int count_0 = 0;
    int temp_1 = 0;
    int temp_0 = 0;
    // 1のみを含む部分配列をカウント
    for (int i = 0; i < size; i++){
        if (arr[i] == 1){
            count_1++;
        }
        else{
            int temp_one_1 = (count_1) * (count_1 + 1) / 2;
            temp_1 = temp_1 + temp_one_1;
            count_1 = 0;
        }
    }
    // 0のみを含む部分配列をカウント
    for (int i = 0; i < size; i++){
        if (arr[i] == 0)
            { count_0++; }
        else{
            int temp_one_0 = (count_0) * (count_0 + 1) / 2;
            temp_0 = temp_0 + temp_one_0;
            count_0 = 0;
        }
    }
    // 配列末尾で連続列が残っている場合の処理
    if (count_1){
        int temp_one_1 = (count_1) * (count_1 + 1) / 2;
        temp_1 = temp_1 + temp_one_1;
    }
    if (count_0){
        int temp_one_0 = (count_0) * (count_0 + 1) / 2;
        temp_0 = temp_0 + temp_one_0;
    }
    cout<<"Subarrays with only 0's are : "<<temp_0;
    cout<<"\nSubarrays with only 1's are : "<<temp_1;
}
int main(){
    int arr[] = { 0, 0, 0, 1, 1, 0, 1};
    int size = sizeof(arr) / sizeof(arr[0]);
    sub_zero_one(arr, size);
    return 0;
}

実行結果

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

Subarrays with only 0's are : 7
Subarrays with only 1's are : 4

まとめ

本記事では、0と1のみで構成されるバイナリ配列において、「0のみ」または「1のみ」を含む部分配列の数を効率的に数える方法を解説しました。連続する同じ要素の区間ごとに n×(n+1)/2 の公式を適用することで、配列を2回走査するだけで答えを求められます。時間計算量は O(N)、空間計算量は O(1) となり、非常に効率的なアプローチです。