P1600 [NOIP2016 提高组] 天天爱跑步 - 树上差分

我只会这题的 \(O(n \log n)\) 线段树合并做法……感觉树上差分的做法好妙。

首先把每条路径 \(s_i\to t_i\) 拆成一条向上的路径和一条向下的路径：\(s_i\to \operatorname{LCA}(s_i,t_i)\to t_i\)。记 \(l_i=\operatorname{LCA}(s_i,t_i)\)。

对于向上的路径，可以发现，如果一条路径 \(s_i\to l_i\) 对观察员 \(u\) 产生了贡献，那么 \(u\in anc_{s_i}\) 且 \(dep_{s_i}=dep_u+w_u\)。

对于向下的路径，推出来 \(dep_{s_i}-2\times dep_{l_i}=w_u-dep_u\)。

向上路径的做法和向下路径的做法类似，下面只讨论向上的做法。

考虑树上差分：将每种数看成一种物品，我们在点 \(s_i\) 打一个 \(+dep_{s_i}\) 的标记，在 \(fa_{l_i}\) 处打一个 \(-dep_{s_i}\) 的标记。

我们记录一个全局的数组 \(c\)，\(c_i\) 表示物品 \(i\) 的个数。

我们对树进行 dfs。当 dfs 到点 \(u\) 时，先记录此时 \(c_{w_u-dep_u}\) 的值 \(cnt\)，再 dfs 点 \(u\) 的所有儿子，然后将点 \(u\) 的所有标记作用到 \(c\) 数组上，此时点 \(u\) 作为观察员的答案就会加上 \(c_{w_u-dep_u}-cnt\)。

为什么这是对的？因为 \(c\) 数组具有可减性，所以我们可以很方便地去除点 \(u\) 的祖先的贡献（通过减去 \(cnt\)）。而且我们在更新点 \(u\) 的答案时，只遍历了点 \(u\) 的所有子孙。

代码

#include <cstdio>
#include <cstring>
#include <cctype>
#include <vector>
using namespace std;
#define For(Ti,Ta,Tb) for(int Ti=(Ta);Ti<=(Tb);++Ti)
#define Dec(Ti,Ta,Tb) for(int Ti=(Ta);Ti>=(Tb);--Ti)
template<typename T> void Read(T &x){
	x=0;int _f=1;
	char ch=getchar();
	while(!isdigit(ch)) _f=(ch=='-'?-1:_f),ch=getchar();
	while(isdigit(ch)) x=x*10+(ch^48),ch=getchar();
	x=x*_f;
}
template<typename T,typename... Args> void Read(T &x,Args& ...others){
	Read(x);Read(others...);
}
typedef long long ll;
const int N=3e5+5;
int n,m,w[N],s[N],lca[N],t[N];
vector<int> G[N];
int dep[N],siz[N],hson[N],fa[N];
void HLD1(int u){
	dep[u]=dep[fa[u]]+1,siz[u]=1;
	for(int v:G[u]){
		if(v==fa[u]) continue;
		fa[v]=u;
		HLD1(v);siz[u]+=siz[v];
		if(!hson[u]||siz[v]>siz[hson[u]]) hson[u]=v;
	}
}
int top[N];
void HLD2(int u,int tp){
	top[u]=tp;
	if(hson[u]) HLD2(hson[u],top[u]);
	else return;
	for(int v:G[u]){
		if(v==fa[u]||v==hson[u]) continue;
		HLD2(v,v);
	}
}
int Lca(int u,int v){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) swap(u,v);
		u=fa[top[u]];
	}
	return dep[u]>dep[v]?v:u;
}
int ans[N],c[N<<1];vector<int> chg1[N];
void Dfs1(int u){
	int cnt=c[w[u]+dep[u]];
	for(int v:G[u]){
		if(v==fa[u]) continue;
		Dfs1(v);
	}
	for(int chg:chg1[u]){
		if(chg>0) ++c[chg];
		else --c[-chg];
	}
	ans[u]+=c[w[u]+dep[u]]-cnt;
}
vector<int> chg2[N];
void Dfs2(int u){
	int cnt=c[w[u]-dep[u]+n];
	for(int v:G[u]){
		if(v==fa[u]) continue;
		Dfs2(v);
	}
	for(int chg:chg2[u]){
		if(chg>0) ++c[chg];
		else --c[-chg];
	}
	ans[u]+=c[w[u]-dep[u]+n]-cnt;
}
int main(){
	Read(n,m);
	int u,v;
	For(i,1,n-1){
		Read(u,v);
		G[u].push_back(v),G[v].push_back(u);
	}
	HLD1(1);HLD2(1,1);
	For(i,1,n) Read(w[i]);
	For(i,1,m){
		Read(s[i],t[i]);lca[i]=Lca(s[i],t[i]);
		chg1[s[i]].push_back(dep[s[i]]);
		if(fa[lca[i]]) chg1[fa[lca[i]]].push_back(-dep[s[i]]);
		chg2[t[i]].push_back(dep[s[i]]-2*dep[lca[i]]+n);
		chg2[lca[i]].push_back(-dep[s[i]]+2*dep[lca[i]]-n);
	}
	Dfs1(1);
	memset(c,0,sizeof c);
	Dfs2(1);
	For(i,1,n) printf("%d ",ans[i]);
	return 0;
}