yukicoder : No.1365 [Cherry 1st Tune] Whose Fault?

これ系はとっさに思い浮かばないな。
https://yukicoder.me/problems/no/1365

問題

N頂点の無向グラフが与えられる。
N頂点は初期状態でエネルギーA[i]を持つ。
辺でつながった頂点間は、総和が変わらないようにエネルギーのやり取りができる。

頂点vの最終的なエネルギーをE[v]、与えられる目標エネルギーをB[v]、与えられる重みをP[v]とする。
この時、不安度はsum(P[v]*(E[v]-B[v])^2)で定義される。
最適なエネルギーのやり取りをしたとき、不安度の最小値を求めよ。

解法

連結成分内では任意にエネルギーを融通できる。
連結成分内のエネルギーの総和をSとすると、sum(E[v])-S=0の範囲でsum(P[v]*(E[v]-B[v])^2)を最小化する問題になる。
ラグランジュの未定乗数法を用いると、sum(P[v]*(E[v]-B[v])^2)の最小値は(sum(E[v])-sum(B[v]))^2 / sum(1/P[v])になる。
そこで、Union Findを用いてE[v],B[v],1/P[v]の和を管理していくとよい。

int N,Q;
double D[202020];
double R[202020];
template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit() {int i; FOR(i,um) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<202020> uf;
double S;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>D[i];
	FOR(i,N) {
		cin>>x;
		D[i]-=x;
	}
	FOR(i,N) {
		cin>>x;
		R[i]=1.0/x;
		S+=D[i]*D[i]/R[i];
	}
	cin>>Q;
	while(Q--) {
		cin>>x>>y;
		x--,y--;
		x=uf[x];
		y=uf[y];
		if(x!=y) {
			S-=D[x]*D[x]/R[x];
			S-=D[y]*D[y]/R[y];
			uf(x,y);
			D[x]=D[y]=D[x]+D[y];
			R[x]=R[y]=R[x]+R[y];
			S+=D[x]*D[x]/R[x];
		}
		_P("%.12lf\n",S);
	}
	
}

まとめ

未定乗数法から先の式変形がさっと出る気がしないなぁ。