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;
}