题解：AT_abc482_e [ABC482E] Diameter Endpoint Colors

题目链接：AtCoder 洛谷

思路

直接做需要对直径长度的奇偶性分讨，比较麻烦，可以把原树的每条边表示成一个新点连接两个点，这样就能保证直径有奇数个点。

根据直径的性质，随便找一条直径，它的中点一定是任意一个直径的中点。把这个中点设为根，依次找它的每个子树下能作为直径端点的叶子，开一个桶存每一种颜色的出现次数，对每个子树的能作为直径端点的叶子，把它们和之前子树中相同颜色的叶子匹配，再扔到桶里去就行。

时间复杂度 \(O(n)\)。

代码

#include <bits/stdc++.h>
using namespace std;
#define pii pair<int,int>
#define ll long long
const int N=4e5+5,P=1e9+7;

int n,dep[N],c[N],fa[N],bk[N],h;
vector<int> e[N],p;

void add(int x,int y)
{
    e[x].push_back(y);
    e[y].push_back(x);
}

void dfs1(int x)
{
    for(int y:e[x])
        if(y!=fa[x])
        {
            fa[y]=x;
            dep[y]=dep[x]+1;
            dfs1(y);
        }
}

void dfs2(int x)
{
    dep[x]=dep[fa[x]]+1;
    if(dep[x]==h)
        p.push_back(c[x]);
    for(int y:e[x])
        if(y!=fa[x])
        {
            fa[y]=x;
            dfs2(y);
        }
}

signed main()
{
    scanf("%d",&n);
    for(int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,n+i);
        add(y,n+i);
    }
    for(int i=1;i<=n;i++)
        scanf("%d",&c[i]);
    if(n==2)
    {
        printf("%d",int(c[1]==c[2]));
        return 0;
    }
    n=n*2-1;
    dfs1(1);
    int x=1;
    for(int i=2;i<=n;i++)
        if(dep[i]>dep[x])
            x=i;
    dep[x]=fa[x]=0;
    dfs1(x);
    x=1;
    for(int i=2;i<=n;i++)
        if(dep[i]>dep[x])
            x=i;
    h=dep[x]/2;
    while(dep[x]>h)
        x=fa[x];
    ll ans=0;
    dep[x]=0;
    for(int y:e[x])
    {
        p.clear();
        fa[y]=x;
        dfs2(y);
        for(int v:p)
            ans+=bk[v];
        for(int v:p)
            bk[v]++;
    }
    printf("%lld",ans);
    return 0;
}