A09【模板】树上前缀和 P4427 [BJOI2018] 求和 - 董晓 - 博客园
P2971 [USACO10HOL] Cow Politics G - 洛谷
给了一颗 n 个节点 边权为 1 的有根树，给了 k 种颜色，树的每个节点已染色。输出同种颜色的节点之间的最远距离。
思路
k 组颜色，询问 k 个最远距离。题目多次询问树上一些路径是权值和，就要考虑树上前缀和了。
每组颜色一定有一个最深的点是最远距离的端点。我们在预处理每个节点深度的时候，顺便记录一下与该点同色的最深点。
我们枚举每个点 x，对应同色的最深点 y，用 dep[x]+dep[y]-2*dep[lca(x,y)] 更新答案。
// 树上前缀和 O(nlogn) #include<bits/stdc++.h> using namespace std;#define N 200005 int h[N],to[N<<1],ne[N<<1],idx; void add(int u,int v){to[++idx]=v,ne[idx]=h[u],h[u]=idx;}int n,k,rt,a[N],p[N]; int fa[N][21],dep[N],mxd[N],point[N],ans[N];void dfs(int x,int f){dep[x]=dep[f]+1; fa[x][0]=f;for(int i=1;i<=20;i++)fa[x][i]=fa[fa[x][i-1]][i-1];if(mxd[a[x]]<dep[x]) mxd[a[x]]=dep[x],point[a[x]]=x; //记录同种颜色最深的点 for(int i=h[x];i;i=ne[i]){int y=to[i];if(y!=f) dfs(y,x);} } int lca(int x,int y){if(dep[x]<dep[y]) swap(x,y); //让x更深for(int i=20;~i;i--)if(dep[fa[x][i]]>=dep[y]) x=fa[x][i]; //x向上跳到y的同一层if(x==y) return x;for(int i=20;~i;i--)if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i]; //一起向上跳return fa[x][0]; } signed main(){scanf("%d%d",&n,&k);for(int i=1;i<=n;++i){scanf("%d%d",&a[i],&p[i]); //颜色,父亲if(p[i]==0) rt=i;else add(i,p[i]),add(p[i],i);}dfs(rt,0); //倍增预处理 dep,fa,point 数组for(int x=1,y;x<=n;++x){y=point[a[x]];ans[a[x]]=max(ans[a[x]],dep[x]+dep[y]-2*dep[lca(x,y)]); //树上前缀和 } for(int i=1;i<=k;++i) printf("%d\n",ans[i]); }