面倒ではあるが、難しくはない。
https://yukicoder.me/problems/no/2654
問題
根付き木が与えられる。
各点には整数値が設定されている。
各点vに対し、以下に答えよ。
根頂点からvのパス上に並ぶ整数値を並べた数列をBとする。
Bが3要素以上の場合、各要素を加減算し、1要素を除きみな等しくしたい。
各要素の変化分の絶対値の総和の最小値を求めよ。
解法
Bに対し、最小値または最大値を孤立させ、残りを一致させてみると良い。
孤立した要素を除くと定番問題で、中央値に合わせるのが最適。
木をDFS探索しながら、Bに含まれる値とその頻度を、座標圧縮したBITを使い管理していこう。
int N; int A[202020],P[202020]; vector<int> As; ll ret[202020]; vector<int> E[202020]; vector<pair<int,int>> V; template<class V, int ME> class BIT { public: V bit[1<<ME],val[1<<ME]; V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;} void add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;} void set(int e,V v) { add(e,v-val[e]);} int lower_bound(V val) { V tv=0; int i,ent=0; for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i); return ent; } }; BIT<ll,20> num,sum; ll hoge(int rem) { num.add(rem,-1); sum.add(rem,-V[rem].first); int tnum=num(N+2); int L,R; if(tnum%2==1) { L=R=num.lower_bound(tnum/2+1); } else { L=num.lower_bound(tnum/2); R=num.lower_bound(tnum/2+1); } ll ret=(sum(N+2)-sum(R-1))-(num(N+2)-num(R-1))*V[R].first; ret+=num(L)*V[R].first-sum(L); if(V[rem].first==V[L].first&&V[rem].first==V[R].first) ret++; num.add(rem,1); sum.add(rem,V[rem].first); return ret; } void dfs(int cur,int pre) { num.add(P[cur],1); sum.add(P[cur],A[cur]); int tnum=num(N+2); if(tnum<=2) { ret[cur]=-1; } else { int L=num.lower_bound(1); int R=num.lower_bound(tnum); if(L==R) { ret[cur]=1; } else { ret[cur]=min(hoge(L),hoge(R)); } } FORR(e,E[cur]) if(e!=pre) dfs(e,cur); num.add(P[cur],-1); sum.add(P[cur],-A[cur]); } void solve() { int i,j,k,l,r,x,y; string s; cin>>N; FOR(i,N+1) { cin>>A[i]; V.push_back({A[i],i}); } sort(ALL(V)); FOR(i,N+1) P[V[i].second]=i; FOR(i,N) { cin>>x>>y; E[x].push_back(y); E[y].push_back(x); } dfs(0,0); FOR(i,N) cout<<ret[i+1]<<endl; }
まとめ
解法は思いついても、こういうのoff-by-oneが怖いなぁ。