C++で最大値と最小値が等しい部分配列の個数を求めるアルゴリズム

はじめに

本記事では、C++を使って「最大要素と最小要素が等しい部分配列」の個数を求める問題を解説します。まずは、問題の具体例から見ていきましょう。

入力 : array = { 2, 3, 6, 6, 2, 4, 4, 4 }
出力 : 12
説明 : {2}, {3}, {6}, {6}, {2}, {4}, {4}, {4}, {6,6}, {4,4}, {4,4}, {4,4,4} が、最大値と最小値が一致する部分配列として成立します。

入力 : array = { 3,3,1,5,1,2,2 }
出力 : 9
説明 : {3}, {3}, {1}, {5}, {1}, {2}, {2}, {3,3}, {2,2} が、最大値と最小値が一致する部分配列として成立します。

解法のアプローチ

例を観察すると、最小値と最大値が等しい部分配列の最小個数は、配列のサイズと一致することが分かります。これは、各要素単体が必ず条件を満たすためです。さらに、連続して同じ数が並んでいる場合は、それらを組み合わせた部分配列が追加で作れるため、個数は増えていきます。

そこで、すべての要素を走査し、直後に続く数が同じかどうかを確認するアプローチが使えます。連続する数が同じであればカウントを増やし、異なる数が見つかった時点で内側のループを抜けます。

そして、内側のループが終了または中断するたびにresult変数へ加算していき、最終的にresult変数の値を答えとして出力します。

コード例

#include <bits/stdc++.h>
using namespace std;
int main(){
    int a[ ] = { 2, 4, 5, 3, 3, 3 };
    int n = sizeof(a) / sizeof(a[0]);
    int result = n, count = 0;
    for (int i = 0; i < n; i++) {
        for (int j = i+1; j < n; j++) {
            if(a[i]==a[j])
                count++;
            else
                break;
        }
        result += count;
        count = 0;
    }
    cout << "Number of subarrays having minimum and maximum elements same:" << result;
    return 0;
}

出力

Number of subarrays having minimum and maximum elements same: 9
Time complexity = O(n<sup>2</sup>)

コードの解説

このコードでは、配列のサイズを格納する変数nを用意し、resultをnで初期化しています。これは、最低でもn個の部分配列（各要素単体）が作れるためです。count変数は、連続する同じ数を数えるために使います。

外側のループは配列の各要素を処理するためのもので、内側のループは現在のインデックス以降にいくつの連続した同じ数があるかを調べます。内側のループが終了するたびに、count変数の値をresult変数に加算し、最後にresult変数に格納された値を出力します。

効率的な解法（O(n)）

このアプローチでも同様にすべての要素を走査しますが、連続して現れる同じ数のかたまりごとに処理を行います。同じ数が見つかるたびにcount変数を増やし、異なる数が見つかった時点で、「n × (n + 1) ÷ 2」という公式を使って、count個の同じ数から作れる部分配列の総数を求め、result変数に加算します。

この公式は、長さnの連続区間から作れる部分配列の総数（n + (n-1) + … + 1 = n(n+1)/2）を表しており、これにより二重ループを排除できます。

コード例（効率版）

#include <bits/stdc++.h>
using namespace std;
int main(){
    int a[] = { 2, 4, 5, 3, 3, 3 };
    int n = sizeof(a) / sizeof(a[0]);
    int result = 0;
    int count = 1, temp = a[0];
    for (int i = 1; i < n; i++) {
        if (temp == a[i]){
            count++;
        }
        else{
            temp = a[i];
            result = result + (count*(count+1)/2);
            count = 1;
        }
    }
    result = result + (count*(count+1)/2);
    cout << "Number of subarrays having minimum and maximum elements same:" << result;
    return 0;
}

出力

Number of subarrays having minimum and maximum elements same: 9
Time complexity : O(n)

コードの解説

このコードでは、まず配列の先頭要素（インデックス0）をtemp変数に格納し、ループはインデックス1から開始します。temp変数と現在の要素が等しければcountを1ずつ増やし、等しくない場合はcount個の同じ数から作れる部分配列の組み合わせ数を公式で計算してresultに加算します。その後、tempを現在の要素に更新し、countを1にリセットします。ループを抜けた後も、末尾の連続区間分の計算が必要なため、再度公式を適用しています。

まとめ

本記事では、最小値と最大値が等しい部分配列の個数を求める問題を取り上げました。シンプルなO(n²)の解法と、連続する同じ数のかたまりごとに計算することで高速化したO(n)の解法の2通りを、C++のコード付きで解説しました。同じロジックはC言語、Java、Pythonなど他の言語でも実装可能です。本記事が皆さんの学習の一助となれば幸いです。