C++14(g++5.4) 解法, 执行用时: 823ms, 内存消耗: 254428K, 提交时间: 2020-04-20 05:42:11

#include<bits/stdc++.h>
#define int long long
#define N 2000010
#define fer(x,y,z) for(register int x=y;x<=z;x++)
using namespace std;
map<int,int>a[N],tag[N];map<int,int>::iterator lt; 
int las[N],to[N],nxt[N],cnt,n,m,rt,tru[N],num[N];
void add(int x,int y){nxt[++cnt]=las[x],las[x]=cnt,to[cnt]=y;}
void sak_dfs(int x,int f,int dep){
	for(int i=las[x];i;i=nxt[i])if(to[i]!=f){
		sak_dfs(to[i],x,dep+1);
		if(a[tru[to[i]]].size()>a[tru[x]].size())swap(tru[x],tru[to[i]]);
		int t1=tru[x],t2=tru[to[i]];
		for(lt=a[t2].begin();lt!=a[t2].end();lt++){
			if(a[t1][lt->first])a[t1][lt->first]=max (a[t1][lt->first] - (tag[t1][lt->first]-dep) ,1ll);
			if(a[t2][lt->first])a[t2][lt->first]=max (a[t2][lt->first] - (tag[t2][lt->first]-dep) ,1ll);
			tag[t1][lt->first]=dep,a[t1][lt->first]+=a[t2][lt->first];
		} 
	}if(num[x])a[tru[x]][dep]+=num[x],tag[tru[x]][dep]=dep;
} 
main(){
	cin>>n>>rt;int x,y;
	fer(i,1,n)scanf("%lld",&num[i]),tru[i]=i;
	fer(i,1,n-1)scanf("%lld%lld",&x,&y),add(x,y),add(y,x);
	sak_dfs(rt,0,0);int ans=0;
	fer(i,0,n)if(a[tru[rt]][i])ans+=max(a[tru[rt]][i]-1-tag[tru[rt]][i],1ll);
	cout<<ans;
}

C++11(clang++ 3.9) 解法, 执行用时: 633ms, 内存消耗: 52572K, 提交时间: 2020-04-19 11:51:18

#include<bits/stdc++.h>
#define LL long long
using namespace std;
const int N=200010;
int n, rt;
map<int, LL> ma[N];
int a[N], dep[N];
vector<int>v[N];
LL tag[N];
void Ins(int u,int d, LL z)
{
	if(!ma[u].count(d))
		ma[u][d] = z + tag[u];
	else
		ma[u][d] = max(ma[u][d], tag[u]+1) + z;
}
void Merge(int u,int v)
{
	if(ma[v].size() > ma[u].size())
	{
		swap(ma[u], ma[v]);
		swap(tag[u], tag[v]);
	}
	for(auto i : ma[v])
		if(i.second)
			Ins(u,i.first, max(i.second-tag[v],1ll));
}
void dfs(int u, int ff)
{
	dep[u] = dep[ff] + 1;
	for(auto to : v[u])
		if(to != ff)
		{
			dfs(to, u);
			Merge(u, to);
		}
	if(a[u])
		Ins(u, dep[u], a[u]);
	tag[u]++;
}
int main()
{
    cin >> n >> rt;
	for(int i = 1; i <= n; i++)
		cin >> a[i];
	for(int i = 1; i < n; i++)
	{
		int x, y;
		cin >> x >> y;
		v[x].push_back(y);
		v[y].push_back(x);
    }
	dfs(rt, 0);
	LL ans = 0;
	for(auto k : ma[rt])
		if(k.second)
			ans += max(1ll, k.second-tag[rt]);
    cout << ans;
}