题解 P5588 【小猪佩奇爬树】

简述题意：树上每个点有一种颜色。对于每种颜色，求经过**所有该颜色的点**的路径数量。 --- 刚看完这题秒出一个$O(n^2)$算法，准备拿60...因为是暴力分就只开了$n=10000$的空间 然后考完发现这个玩意跑得挺快的 然后开大空间，开`long long`交一发（ **过了。。。** 我\*\*\*\*\* --- 首先我们可以想到，对于一种颜色的点，如果存在这样的路径，则只有两种情况： * 存在一个深度最大的点，其他这种颜色的点都是其祖先，形成一条链。 * 存在两条如上一条所述且互不相交的链 如果都不是，直接判无解就行。 而对于不存在的颜色，答案显然是$\frac{n(n-1)}{2}$ 然后上面这两种情况都会得到一条经过所有该颜色的点的最短路径。 把链上面的树枝全部砍掉，只剩下两个端点$u,v​$上的子树 我们发现答案就是两个子树大小乘起来。 交一发，WA 0 pts... 我们可以发现这条链可能是一个点。那么把以该点为根时所有子树的大小两两相乘就是答案。。。 然后就可以$O(n^2)$\*过去了。。。 我找链用了set+暴力跳父亲，因此复杂度是错的...不知道正解该怎么做... ```cpp #include<cstdio> #include<set> #define int long long const int N=1<<15; int n,head[N],size[N],fa[N],dep[N],col[N]; struct Edge { int next,to; }; Edge E[N<<1]; void add(int u,int v) { static int tot=0; E[++tot].next=head[u]; E[tot].to=v; head[u]=tot; } std::set<std::pair<int,int> > pts[N]; void dfs(int u) { size[u]=1; pts[col[u]].insert(std::make_pair(-dep[u],u)); for(int i=head[u];i;i=E[i].next) { int v=E[i].to; if(v==fa[u])continue; fa[v]=u; dep[v]=u+1; dfs(v); size[u]+=size[v]; } } int check_ans(int c) { if(pts[c].empty())return n*(n-1)/2; if(pts[c].size()==1) { int u=(*pts[c].begin()).second; int sum=0,res=0; for(int i=head[u];i;i=E[i].next) { int v=E[i].to; if(v==fa[u])continue; res+=sum*size[v]; sum+=size[v]; } if(fa[u]) { res+=(n-size[u])*sum; } return res+n-1; } int u=(*pts[c].begin()).second,v,t=u; pts[c].erase(std::make_pair(-dep[u],u)); int except=u; while(fa[t]) { if(col[fa[t]]==c) { v=fa[t]; except=t; pts[c].erase(std::make_pair(-dep[fa[t]],fa[t])); } t=fa[t]; } if(pts[c].empty()) { return size[u]*(n-size[except]); } v=(*pts[c].begin()).second,t=v; pts[c].erase(std::make_pair(-dep[v],v)); while(fa[t]) { if(col[fa[t]]==c) { if(pts[c].erase(std::make_pair(-dep[fa[t]],fa[t]))==0)return 0; } t=fa[t]; } if(!pts[c].empty())return 0; return size[u]*size[v]; } signed main() { scanf("%lld",&n); for(int i=1;i<=n;++i)scanf("%lld",col+i); for(int i=1;i<n;++i) { int u,v; scanf("%lld%lld",&u,&v); add(u,v); add(v,u); } dfs(1); for(int i=1;i<=n;++i)printf("%lld\n",check_ans(i)); } ```