Official

D - ケーキの均等分配 / Equal Distribution of Cake Editorial by physics0523


ケーキの総和 \(s\)\(N\) の倍数でないとき均等に分けることは不可能なので、そのようなケースを除去して考えます。

各ケーキについて、現在誰が持っているかを \(x\) 、最終的に誰が持っているかを \(y\) とすると、この問題は全てのケーキについて \(|x-y|\) の総和を最小化する問題になります。

すると、以下を満たす戦略が最善であることがわかります。

  • 子ども \(x_1\) が持つあるケーキの行き先が子ども \(y_1\) であるとする。
  • 子ども \(x_2\) が持つあるケーキの行き先が子ども \(y_2\) であるとする。
  • このとき、任意のケーキの組について \(x_1 < x_2\) ならば \(y_1 \le y_2\) が成り立つべきである。

これに反するような戦略を取った場合、当該ケーキの組の行き先を入れ替えても損しません。

ここから、以下の戦略が導かれます。

  • 最終的に各子どもは \(s'=s/N\) 個のケーキを持っているべきである。
  • 子ども \(i=1,2,\dots,N\) について、順に以下を行う。
    • 子ども \(i\) が持っているケーキのうち行き先が未定のものがなくなるまで、以下を繰り返す。
      • 現在ケーキが足りていない子どものうち最も番号が若い子ども \(j\) に対し、ケーキが足りるか子ども \(i\) の行き先未定なケーキが無くなるまで、ケーキの行き先を子ども \(j\) に割り当てる。

端的に言うと、子どもの番号の若い順にケーキの行き先を若い方から埋めていく感じです。

この戦略をケーキごとにやっていくと実行時間制限に間に合いませんが、結果が変わらない間はまとめて移動を行うことにすると全体で時間計算量 \(O(N)\) の解法が得られます。

実装例 (C++):

#include<bits/stdc++.h>

using namespace std;
using ll=long long;

int main(){
  ll N;
  cin >> N;
  vector<ll> A(N);
  ll s=0;
  for(auto &nx : A){
    cin >> nx;
    s+=nx;
  }
  if(s%N){cout << "-1\n"; return 0;}
  s/=N;

  ll res=0;
  ll pos=0,cnt=0;
  for(ll i=0;i<N;i++){
    while(A[i]>0){
      ll del=min(A[i],s-cnt);
      A[i]-=del;
      cnt+=del;
      res+=abs(pos-i)*del;
      if(cnt==s){
        pos++;
        cnt=0;
      }
    }
  }
  cout << res << "\n";
  return 0;
}

posted:
last update: