CF1654G Snowy Mountain

题面传送门
这步根号是实在想不到。
考虑我们最优策略是什么，显然我们目的是消耗最多的动能，容易发现如果一对点直接有边并且高度相等，那么这对点是可以消耗掉目前为止所有动能的。
那么目的就是找到一个点能走到的点中高度最低的等高点\(f_i\)，那么答案就是\(2h_i-f_i\)。
正着做不太好做，考虑反过来，让每一个高度相同点去找所有能到达它的点。
然后最神仙一步就来了：这种高度相同点的本质不同高度数只有\(O(\sqrt n)\)个。
如果存在一对高度相同点，那么这对点实际上各自需要至少\(h_i\)个点铺垫，那么所有高度和应该为\(O(n)\)，故本质不同高度数为\(O(\sqrt n)\)的。
然后对于每组高度相同点直接多源bfs就好了。时间复杂度\(O(n\sqrt n)\)
code:

#include<bits/stdc++.h>
#define I inline
#define max(a,b) ((a)>(b)?(a):(b))
#define min(a,b) ((a)<(b)?(a):(b))
#define abs(x) ((x)>0?(x):-(x))
#define re register
#define RI re int
#define ll long long
#define db double
#define lb long db
#define N (200000+5)
#define M ((1<<19)+5)
#define mod 998244353
#define Mod (mod-1)
#define eps (1e-9)
#define U unsigned int
#define it iterator
#define Gc() getchar() 
#define Me(x,y) memset(x,y,sizeof(x))
#define Mc(x,y) memcpy(x,y,sizeof(x))
#define d(x,y) (n*(x-1)+(y))
#define R(n) (rand()*rand()%(n)+1)
#define Pc(x) putchar(x)
#define LB lower_bound
#define UB upper_bound
#define PB push_back
using namespace std;vector<int> S[N],P[N];
int n,m,k,H[N],F[N],x,y,G[N];queue<int> Q;
int main(){
	freopen("1.in","r",stdin);
	Me(H,0x3f);RI i,j;scanf("%d",&n);for(i=1;i<=n;i++)scanf("%d",&x),x&&(H[i]=0,Q.push(i),0);for(i=1;i<n;i++)scanf("%d%d",&x,&y),S[x].PB(y),S[y].PB(x);
	while(!Q.empty()){x=Q.front();Q.pop();for(RI i:S[x]) H[i]>H[x]+1&&(H[i]=H[x]+1,Q.push(i),0);};for(i=1;i<=n;i++) for(RI j:S[i]) if(H[i]==H[j]){P[H[i]].PB(i);break;}
	Mc(F,H);for(i=0;i<=n;i++) {
		if(!P[i].size()) continue;Me(G,-0x3f);for(RI j:P[i]) G[j]=0,Q.push(j);
		while(!Q.empty()) {x=Q.front();Q.pop();for(RI j:S[x]) H[j]>=H[x]&&G[j]<min(G[x],0)+(H[j]>H[x]?1:-1)&&(G[j]=min(G[x],0)+(H[j]>H[x]?1:-1),Q.push(j),0);};
		for(j=1;j<=n;j++) G[j]>=0&&(F[j]=min(F[j],i));
	}for(i=1;i<=n;i++) printf("%d ",2*H[i]-F[i]);
}