【C++】d日以内に最初の山へ集められる干し草の俵の最大数を求めるアルゴリズム

問題概要

n個の要素を持つ配列Aと値dが与えられます。ある農家では、n個の干し草の山が一列に並べられており、i番目の山にはA[i]個の干し草が積まれています。

ここに1頭の牛がいて、1日に1回だけ、どれかの山から隣接する山へ干し草を1つ移動させることができます(何もしない日があっても構いません)。牛はd日間で1番目の山の干し草の数をできるだけ多くしたいと考えています。d日後の1番目の山にある干し草の最大数を求めてください。

例えば、入力が d = 5、A = [1, 0, 3, 2] の場合、出力は 3 になります。これは次のような手順で移動できるためです。

• 1日目：3番目の山 → 2番目の山へ1つ移動
• 2日目：3番目の山 → 2番目の山へさらに1つ移動
• 3日目：2番目の山 → 1番目の山へ1つ移動
• 4日目：2番目の山 → 1番目の山へさらに1つ移動

こうして、1番目の山は 1 + 2 = 3 個の干し草を持つことになります。

解法のアプローチ

この問題のポイントは、i番目の山から1番目の山へ干し草を1つ運ぶのにi回の移動(つまりi日)が必要だという点です。なぜなら、牛は隣接する山にしか移動できないため、距離iだけ離れた山からの輸送にはi日かかるからです。

したがって、残り日数dの中でi番目の山から運べる干し草の数は、以下のように表せます。

• w = min(A[i], d / i)

つまり「その山に元々ある干し草の数」と「残り日数で運べる上限」の小さい方を採用し、それをa0に加算しながら、使用した日数分(w × i)をdから差し引いていきます。これをすべての山に対して繰り返すことで答えが得られます。

アルゴリズムの手順

a0 := A[0]
n := size of A
for initialize i := 1, when i < n, update (increase i by 1), do:
    ai := A[i]
    w := minimum of ai and d / i
    a0 := a0 + w
    d := d - w * i
return a0

C++での実装例

以下に実際のC++コードを示します。

#include <bits/stdc++.h>
using namespace std;
int solve(int d, vector<int> A){
    int a0 = A[0];
    int n = A.size();
    for (int i = 1; i < n; i++){
        int ai = A[i];
        int w = min(ai, d / i);
        a0 += w;
        d -= w * i;
    }
    return a0;
}
int main(){
    int d = 5;
    vector<int> A = { 1, 0, 3, 2 };
    cout << solve(d, A) << endl;
}

入力

5, { 1, 0, 3, 2 }

出力

3

計算量について

このアルゴリズムは配列を一度走査するだけで済むため、時間計算量はO(n)、空間計算量はO(1)と非常に効率的です。貪欲法(Greedy法)の典型的な応用例であり、「遠い山ほどコストが高い」という性質を利用して、各山から独立に最適値を決定しています。