Official

G - Approximate Equalization Editorial by mechanicalpenciI


まず、最初に、整数列 \(A=(A_1,A_2,\ldots,A_N)\) に加えて目標とする数列 \(B=(B_1,B_2,\ldots,B_N)\) が与えられた時に、 \(A\)\(B\) に変えるために必要な操作回数の最小値について考えます。

まず、総和が等しくない、すなわち \(\displaystyle\sum_{i=1}^N A_i \neq \sum_{i=1}^N B_i\) ならば \(A\)\(B\) に変える事はできません。なぜなら、問題文中の操作によって総和が変化することはないためです。
一方、総和が等しい時、つねに \(A\)\(B\) に変える事ができ、\(0\leq i\leq N\) について、\(S_i=\displaystyle\sum_{k=1}^i A_k\), \(T_i=\displaystyle\sum_{k=1}^i B_k\)とおいたときの、\(\displaystyle\sum_{i=1}^{N-1}\lvert S_i-T_i\rvert\) の値に等しいです。 なぜなら、問題文中の操作を \(S_i\) についての操作に置き換えると、

  • \(1\leq i\leq N-1\) をみたす整数 \(i\) を選び、\(S_i\)\(1\) 減らす。
  • \(1\leq i\leq N-1\) をみたす整数 \(i\) を選び、\(S_i\)\(1\) 増やす。

であり、\((S_0=0,S_1,S_2,\ldots,S_{N-1},S_N)\)\((A_1,A_2,\ldots,A_N)\) および \((T_0=0,T_1,T_2,\ldots,T_{N-1},T_N)\)\((B_1,B_2,\ldots,B_N)\) は一対一に対応しているからです。(総和が等しいという条件より、\(S_N=T_N\) となっていることに注意してください。)

さて、次に \(A\) が与えられた時に条件をみたす最終的な整数列 \(B=(B_1,B_2,\ldots,B_N)\) の候補について考えます。まず、問題文中の条件から、ある整数 \(x\) について、各 \(B_i\)\(x\) または \(x+1\) である必要があります。さらに、\(A\) から操作によって変えられる数列の条件から、\(\displaystyle\sum_{i=1}^N B_i=S_N\) である必要があります。

この時、ある \(i\) について \(B_i\leq \frac{S_N}{N}-1\) ならば、\(\displaystyle\sum_{i=1}^N B_i\leq (\frac{S_N}{N}-1)+(N-1)\times\frac{S_N}{N}=S_N-1\) より不適となります。 同様に、ある \(i\) について \(B_i\geq \frac{S_N}{N}+1\) ならば、\(\displaystyle\sum_{i=1}^N B_i\geq (\frac{S_N}{N}+1)+(N-1)\times\frac{S_N}{N}=S_N+1\) より不適となります。 よって、各 \(B_i\)\(\frac{S_N}{N}-1<B_i<\frac{S_N}{N}+1\) をみたす整数である必要があります。

  • \(S_N\)\(N\) の倍数でない時

任意の \(i\) について \(B_i=\lfloor \frac{S_N}{N}\rfloor\) または\(\lfloor\frac{S_N}{N}\rfloor+1\) である必要があります。 (ここで、\(\lfloor\frac{S_N}{N}\rfloor\)\(\frac{S_N}{N}\) 以下の最大の整数を表します。) \(1\leq i\leq N\) をみたす整数 \(i\) のうち、 \(B_i=\lfloor\frac{S_N}{N}\rfloor+1\)であるような \(i\)\(k\) 個、 \(B_i=\lfloor\frac{S_N}{N}\rfloor\)であるような \(i\)\((N-k)\) 個とすると、\(S_N=\displaystyle\sum_{i=1}^NB_i=N\cdot\left\lfloor\frac{S_N}{N}\right\rfloor+k\equiv k\pmod{N}\) となります。よって、\(S_N=N\cdot \lfloor\frac{S_N}{N}\rfloor+r\) (\(1\leq r\leq N-1)\) と表した時、\(B\)\(r\) 個の \(\lfloor\frac{S_N}{N}\rfloor+1\)\((N-r)\) 個の \(\lfloor\frac{S_N}{N}\rfloor\) からなる必要があり、逆にこれをみたすならば問題文中の条件をみたし、さらに \(A\) から \(B\) に変える事が可能です。

  • \(S_N\)\(N\) の倍数の時

任意の \(i\) について \(B_i=\frac{S_N}{N}\) である必要があり、実際このとき条件をみたします。これは\(S_N\)\(N\) の倍数でない時の \(r=0\) の場合として考える事ができます。

\(B\) の候補全体についての\(\displaystyle\sum_{i=1}^{N-1}\lvert S_i-T_i\rvert\) の最小値は動的計画法によって求める事ができます。

\(dp[i][j]\) \((0\leq i\leq N, 0\leq j\leq \min(i,r))\) を、\(B_1,B_2,\ldots,B_i\) のうち \(B_k=\lfloor\frac{S_N}{N}\rfloor+1\) であるような \(k\)\(j\) 個ある時の \(\displaystyle\sum_{k=1}^{i}\lvert S_k-T_k\rvert\) の値とすると、\(dp[0][0]=0\) であり、

\[ dp[i][j]=\min(dp[i-1][j-1],dp[i-1][j])+\left\lvert i\cdot\left\lfloor\frac{S_N}{N}\right\rfloor +j-S_i\right \rvert \]

(ただし、\(j=0\) ならば \(dp[i-1][j-1]=+\infty\), \(j=i\)ならば \(dp[i-1][j]=+\infty\) ) として計算する事ができます。 答えは \(dp[N][r]\) の値となります。 各 \(dp[i][j]\)\(dp[i-1][j-1]\) および \(dp[i-1][j]\) から \(O(1)\) で計算する事ができるため、 全体で \(O(N^2)\) であり、制約は \(N\leq 5000\) である事から、十分高速に答えを求める事ができました。

c++ による実装例:

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

int main(void){
	int n,r;
	long long d;
	long long a[5000],dp[5000],dp2[5000];

	cin>>n;
	for(int i=0;i<n;i++)cin>>a[i];

	for(int i=0;i<n-1;i++)a[i+1]+=a[i];
	if(a[n-1]>=0){
		d=a[n-1]/n;
		r=a[n-1]%n;
	}
	else{
		r=(n-(-a[n-1])%n)%n;
		d=(a[n-1]-r)/n;
	}

	dp[0]=0;
	for(int i=0;i<n;i++){
		dp2[0]=dp[0]+abs(d*(i+1)-a[i]);
		for(int j=0;j<min(r,i);j++)dp2[j+1]=min(dp[j],dp[j+1])+abs(d*(i+1)+j+1-a[i]);
		if(i+1<=r)dp2[i+1]=dp[i]+abs((d+1)*(i+1)-a[i]);
		for(int j=0;j<=min(r,i+1);j++)dp[j]=dp2[j];
	}

	cout<<dp[r]<<endl;
	return 0;
}

posted:
last update: