【C++解説】最大k回の加算操作で等しくできる要素数の最大化アルゴリズム

この問題では、与えられた配列に対して、各要素を合計で最大k回まで増加（+1）できるという条件下で、等しい値にできる要素の最大個数を求めます。

まず、具体例を使って何をすべきか確認してみましょう。

入力例1

a[] = {1, 3, 8}, k = 4

出力例1

2

説明

この例では、要素「1」を3回、「3」を4回（計7回…と見せかけて、実際は各要素への加算回数の合計がk以内に収まるように調整）インクリメントすることで、2つの「4」を作り出せます。結果として配列は a[] = {4, 4, 8} となり、答えは 2 です。

入力例2

arr = {2, 5, 9}, k = 2

出力例2

0

この場合、合計2回の加算ではどの要素も同じ値に揃えられないため、答えは 0 になります。

アルゴリズムのアプローチ

以下のプログラムでは、二分探索と累積和を組み合わせて効率的に答えを求めています。手順は次のとおりです。

• main() 関数内で、配列の要素を格納する int a[]、配列のサイズ size、そして最大更新可能回数 k を初期化します。

• Max() 関数では、まず配列を昇順にソートします。その後、累積和を格納する配列 p[size + 1] と、その位置までの最大値を格納する配列 m[size + 1] を宣言します。

• i = 0 から i <= size までループし、p[i] = 0 および m[i] = 0 で初期化します。

• ループの外側で m[0] = arr[0]、p[0] = arr[0] を設定します。

• i = 1 から i < size までループし、p[i] = p[i - 1] + arr[i] で累積和を計算し、m[i] = max(m[i - 1], arr[i]) でその位置までの最大値を求めます。

• ループ後、左端・右端・最終的な答えを格納するために int Lt = 1, Rt = size, result を初期化し、二分探索を開始します。

• 条件 (Lt < Rt) を持つwhileループを実行します。ループ内で int mid = (Lt + Rt) / 2 を計算し、EleCal(p, m, mid - 1, k, size) が真であれば、result = mid として Lt = mid + 1 を設定します。

• そうでなければ、単純に Rt = mid - 1 とします。

• ループを抜けたら result を出力します。

• 関数 bool EleCal() 内では、条件 for (int i = 0, j = x; j <= size; j++, i++) のforループを開始します。

• ループ内で (x * m[j] - (p[j] - p[i]) <= k) を判定します。真であれば true を返します。ループを抜けたら false を返します。

この判定式のポイントは、ある区間の x 個の要素すべてを、その区間の最大値 m[j] に揃えるために必要な加算回数が x * m[j] - (区間の総和) で表せることです。これが k 以内であれば、x 個の要素を等しくできることを意味します。

実装例

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

// x個の要素を等しくできるかどうかを判定する関数
bool EleCal(int p[], int m[], int x, int k, int size){
    for (int i = 0, j = x; j <= size; j++, i++){
        // 区間の最大値に揃えるのに必要な加算回数がk以下なら可能
        if (x * m[j] - (p[j] - p[i]) <= k)
            return true;
    }
    return false;
}

void Max(int arr[], int size, int k){
    // 配列を昇順にソート
    sort(arr, arr + size);
    
    int p[size + 1];
    // 最大値用の配列
    int m[size + 1];
    
    // 累積和配列と最大値配列を初期化
    for (int i = 0; i <= size; ++i){
        p[i] = 0;
        m[i] = 0;
    }
    
    m[0] = arr[0];
    p[0] = arr[0];
    
    for (int i = 1; i < size; i++){
        // 配列の累積和を計算
        p[i] = p[i - 1] + arr[i];
        // その位置までの最大値を計算
        m[i] = max(m[i - 1], arr[i]);
    }
    
    // 二分探索
    int Lt = 1, Rt = size, result;
    while (Lt < Rt){
        int mid = (Lt + Rt) / 2;
        if (EleCal(p, m, mid - 1, k, size)){
            result = mid;
            Lt = mid + 1;
        }
        else
            Rt = mid - 1;
    }
    
    // 答えを出力
    cout << result;
}

// main関数
int main(){
    int a[] = { 1, 3, 8 };
    int size = sizeof(a) / sizeof(a[0]);
    int k = 4;
    Max(a, size, k);
    return 0;
}

出力

2

計算量について

ソートに O(N log N)、二分探索の各ステップでの判定に O(N) かかるため、全体の時間計算量は O(N log N) となります。素朴な全探索（O(N²) 以上）に比べて大幅に高速化でき、大きな入力サイズにも対応できます。