P3258 [JLOI2014] 松鼠的新家

P3258 [JLOI2014] 松鼠的新家

大意

给多次修改，\(u, v\)，每次将 \(u \to v\) 的路径上的点权加 \(1\)，求最终每个点的点权值。

思路

显然树上差分。

定义 \(d_i\) 表示从 \(i\) 到根的路径上点权和，那么 \(u \to v\) 就只需要 \(d_u\) 和 \(d_v\) 处加 \(1\)，然后在 \(d_{lca(u, v)}\) 和 \(d_{fa_{lca(u,v)}}\) 处减 \(1\)，最后将差分数组求前缀和就是原数组。

代码

#include<iostream>
using namespace std;

const int MAXN = 3 * 1e5 + 5;
struct node{
	int to,next;
}e[MAXN * 2];

int f[MAXN][20],dp[MAXN];
int a[MAXN];
int h[MAXN * 2],tot = 0;
int n,k,ans = -1;
int diff[MAXN];

void add(int x,int y){
	e[++ tot] = {y,h[x]};
	h[x] = tot;
}
void dfs(int u,int fa){
	dp[u] = dp[fa] + 1;
	f[u][0] = fa;
	for(int i = 1;(1 << i) <= dp[u];i ++){
		f[u][i] = f[f[u][i - 1]][i - 1];
	}
	for(int i = h[u];i;i = e[i].next){
		int v = e[i].to;
		if(v != fa) dfs(v,u);
	}
	return;
}

void dfs2(int u,int fa){
	for(int i = h[u];i;i = e[i].next){
		int v = e[i].to;
		if(v == fa) continue;
		dfs2(v,u);
		diff[u] += diff[v];
	}
}

int LCA(int x,int y){
	if(dp[x] < dp[y]) swap(x,y);
	for(int i = 19;i >= 0;i --){
		if(dp[x] - (1 << i) >= dp[y]){
			x = f[x][i];
		}
	}
	if(x == y) return x;
	for(int i = 19;i >= 0;i --){
		if(f[x][i] != f[y][i]){
			x = f[x][i],y = f[y][i];
		}
	}
	return f[x][0];
}
int main(){
	cin >> n;
	for(int i = 1;i <= n;i ++) cin >> a[i];
	for(int u,v,i = 1;i <= n - 1;i ++){
		cin >> u >> v;
		add(u,v); add(v,u);
	}
	dfs(1,0);
	for(int u,v,i = 2;i <= n;i ++){
		u = a[i - 1],v = a[i];
		int l = LCA(u,v);
		diff[u] ++;
		diff[v] ++;
		diff[l] --;
		diff[f[l][0]] --;
	}
	dfs2(1,0);
	for(int i = 2;i <= n;i ++){
		diff[a[i]] --;
	}
	for(int i = 1;i <= n;i ++){
		cout << diff[i] << '\n';
	}
	return 0;
}