P1600 天天爱跑步

P1600 天天爱跑步

题目描述

详见：P1600 天天爱跑步

Solution

树上差分+LCA。

Code

#include<bits/stdc++.h>
using namespace std;
const int MAXN=4e5+50;
vector<int> e[MAXN];
struct wnode{int opt,x,c; };
vector<wnode> w[MAXN];
int dep[MAXN],fa[MAXN][21],Log[MAXN],cnt[2][MAXN],W[MAXN],ans[MAXN],P,n,m;
inline int read()
{
    int x=0,f=1; char c=getchar();
    while (c<'0'||c>'9') { if(c=='-') f=-1; c=getchar(); }
    while (c>='0'&&c<='9') { x=(x<<3)+(x<<1)+(c^48); c=getchar(); }
    return x*f;
}
void dfs(int x,int father)
{
	dep[x]=dep[father]+1;
	fa[x][0]=father;
	for (int i=1;i<=Log[dep[x]];i++) fa[x][i]=fa[fa[x][i-1]][i-1];
	for (int i=0;i<e[x].size();i++)
	    if (e[x][i]!=father) dfs(e[x][i],x);
}
int get_lca(int x,int y)
{
	if (dep[x]<dep[y]) swap(x,y);
	for (int i=Log[dep[x]];i>=0;i--)
	    if (dep[fa[x][i]]>=dep[y]) x=fa[x][i];
	if (x==y) return x;
	for (int i=Log[dep[x]];i>=0;i--)
	    if (fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
	return fa[x][0];
}
void solve(int x,int father)
{
	ans[x]=cnt[0][W[x]+dep[x]]+cnt[1][dep[x]-W[x]+P];
	for (int i=0;i<w[x].size();i++) cnt[w[x][i].opt][w[x][i].x]+=w[x][i].c;
	for (int i=0;i<e[x].size();i++)
	    if (e[x][i]!=father) solve(e[x][i],x);
	ans[x]=cnt[0][W[x]+dep[x]]+cnt[1][dep[x]-W[x]+P]-ans[x];
}
int main()
{
	n=read(),m=read(),P=n*3;
	for (int i=1;i<n;i++) 
	{
		int u=read(),v=read();
		e[u].push_back(v);
		e[v].push_back(u);
	}
	for (int i=1;i<=n;i++) W[i]=read(); 
	dep[0]=-1,Log[1]=0;
	for (int i=2;i<=n;i++) Log[i]=Log[i>>1]+1;
	dfs(1,0);
	for (int i=1;i<=m;i++)
	{
		int u=read(),v=read(),lca=get_lca(u,v),len=dep[u]+dep[v]-dep[lca]*2;
		w[u].push_back((wnode){0,dep[u],1});
		w[v].push_back((wnode){1,dep[v]-len+P,1});
		w[lca].push_back((wnode){0,dep[u],-1});
		w[fa[lca][0]].push_back((wnode){1,dep[v]-len+P,-1});
	}
	solve(1,0);
	for (int i=1;i<=n;i++) printf("%d ",ans[i]);
	return 0;
}