Diameter Endpoint Colors

Diameter Endpoint Colors

树的直径

题目大意

给定一棵有 \(N\) 个顶点的树，每个顶点有一个颜色 \(C_v\)。
求满足以下条件的整数对 \((u,v)\)，\(1\le u<v\le N\) 的数量：

1. \(C_u=C_v\)，即两点颜色相同；
2. \(\operatorname{dist}(u,v)=D\)，其中 \(D\) 是树的直径。

路径长度指边数。

---

核心性质

树的所有直径共享同一个中心：

• 若直径长度 \(D\) 为偶数，设 \(D=2R\)，则中心是唯一一个顶点 \(c\)，所有直径都经过 \(c\)；
• 若直径长度 \(D\) 为奇数，设 \(D=2R+1\)，则中心是唯一一条边 \((c_1,c_2)\)，所有直径都包含这条边。

因此，任何距离等于直径的点对，都必须经过树的中心。

---

算法步骤

1. 求直径和中心

先从任意点（如 \(1\)）BFS/DFS，找到最远点 \(A\)。
再从 \(A\) 出发 BFS/DFS，找到最远点 \(B\)，同时得到每个点到 \(A\) 的距离。

设直径长度为：

\[D=\operatorname{dist}(A,B) \]

再从 \(B\) 出发 BFS/DFS，得到每个点到 \(B\) 的距离。

根据 \(D\) 的奇偶性分类处理。

---

情况一：直径长度为偶数

设：

\[D=2R \]

中心是唯一一个顶点 \(c\)，满足：

\[\operatorname{dist}(A,c)=R,\qquad \operatorname{dist}(B,c)=R \]

距离中心 \(c\) 恰好为 \(R\) 的点，称为最远层点。

若两个最远层点 \(u,v\) 满足：

\[\operatorname{dist}(u,v)=2R \]

当且仅当它们来自中心 \(c\) 的不同邻接分支。

因此问题转化为：

在所有距离中心 \(c\) 恰好为 \(R\) 的点中，统计颜色相同且来自不同分支的点对数量。

容斥统计

设所有最远层点中，颜色 \(color\) 的出现次数为 \(total[color]\)。
设某个分支内，颜色 \(color\) 的出现次数为 \(branch[color]\)。

先计算所有最远层点中同色点对数：

\[\sum_{color}\binom{total[color]}{2} \]

再减去每个分支内部的同色点对数：

\[\sum_{\text{分支}}\sum_{color}\binom{branch[color]}{2} \]

最终答案：

\[ans= \sum_{color}\binom{total[color]}{2} - \sum_{\text{分支}}\sum_{color}\binom{branch[color]}{2} \]

实现时不需要二维数组。
对每个分支单独用一个临时计数数组，统计完该分支后把对应颜色清零即可。

情况二：直径长度为奇数

设：

\[D=2R+1 \]

中心是一条边 \((c_1,c_2)\)，满足：

\[\operatorname{dist}(A,c_1)=R,\qquad \operatorname{dist}(B,c_1)=R+1 \]

\[\operatorname{dist}(A,c_2)=R+1,\qquad \operatorname{dist}(B,c_2)=R \]

删除中心边后，树分成两侧。

若两个点 \(u,v\) 满足：

\[\operatorname{dist}(u,v)=2R+1 \]

则它们必须分别在中心边两侧，并且：

\[\operatorname{dist}(u,c_1)=R,\qquad \operatorname{dist}(v,c_2)=R \]

或者反过来。

因此：

• 在 \(c_1\) 一侧，收集距离 \(c_1\) 恰好为 \(R\) 的点；
• 在 \(c_2\) 一侧，收集距离 \(c_2\) 恰好为 \(R\) 的点；
• 统计左右两侧颜色相同的配对数。

设左侧颜色 \(color\) 出现次数为 \(L[color]\)，右侧为 \(Rcnt[color]\)，则：

\[ans=\sum_{color} L[color]\times Rcnt[color] \]

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pli = pair<ll, int>;
using tpi = tuple<int, int, int>;
const ll INF = 1e18;
const ll MOD = 998244353;
const int N = 2e5 + 10;

// 推理可知不存在平行直径
// 无论奇偶，平行直径都会产生新的更长的直径
// 必定找到一个中心点

void solve() {
    int n; cin >> n;
    vector<vector<int>> g(n + 1, vector<int>());
    for(int i = 1; i < n; i ++) {
        int a, b; cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    vector<int> col(n + 1);
    for(int i = 1; i <= n; i ++) cin >> col[i];
    
    auto bfs = [&](int x) {
        vector<int> d(n + 1, 0);
        vector<int> vis(n + 1, 0); vis[x] = 1;
        queue<int> q; q.push(x);
        while(!q.empty()) {
            int u = q.front(); q.pop();
            for(int v : g[u]) {
                if(vis[v]) continue;
                d[v] = d[u] + 1;
                q.push(v); vis[v] = 1;
            }
        }
        return d;
    };
    vector<int> predist = bfs(1);
    int pointA, mx = 0;
    for(int i = 1; i <= n; i ++)
        if(predist[i] > mx) mx = predist[i], pointA = i;

    vector<int> A_dist = bfs(pointA);
    int pointB, dir = 0; //另一个端点和直径
    for(int i = 1; i <= n; i ++)
        if(A_dist[i] > dir) dir = A_dist[i], pointB = i;
    vector<int> B_dist = bfs(pointB);

    if(dir % 2 == 0) { // 直径为偶，点数为奇，存在唯一中心点
        int center;
        int r = dir / 2;
        for(int i = 1; i <= n; i ++) {
            if(A_dist[i] == r && B_dist[i] == r) {center = i; break;}
        }
        // 可以跟不同子树的相同颜色组成直径
        // 容斥原理减去子树内部即可
        vector<int> dist = bfs(center);
        int len = g[center].size();
        vector<vector<int>> branch(len, vector<int>());
        vector<int> vis(n + 1, 0); vis[center] = 1;
        for(int i = 0; i < len; i ++) { //计算所有子树
            int u = g[center][i]; vis[u] = 1;
            queue<int> q; q.push(u);
            while(!q.empty()) {
                int v = q.front(); q.pop();
                if(dist[v] == r) {branch[i].push_back(col[v]); continue;}
                for(int e : g[v]) {
                    if(vis[e]) continue;
                    q.push(e); vis[e] = 1;
                }
            }
        }

        ll ans = 0;
        vector<int> cnt(N + 1, 0); //临时计数
        vector<int> tol(n + 1, 0);
        for(int i = 0; i < len; i ++) {
            for(int c : branch[i]) cnt[c] ++;
            for(int c : branch[i]) {
                if(cnt[c] > 0) {
                    tol[c] += cnt[c];
                    ans -= (ll)cnt[c] * (cnt[c] - 1) / 2;
                    cnt[c] = 0;
                }
            }
        }
        for(int i = 1; i <= n; i ++) ans += (ll)tol[i] * (tol[i] - 1) / 2;
        cout << ans << "\n";
    }
    else { //奇数情况
        int cenA, cenB;
        int r = dir / 2;
        for(int i = 1; i <= n; i ++) {
            if(A_dist[i] == r && B_dist[i] == r + 1) cenA = i;
            if(A_dist[i] == r + 1 && B_dist[i] == r) cenB = i;
        }

        vector<int> distA = bfs(cenA);
        vector<int> distB = bfs(cenB);
        vector<int> colA(n + 1, 0), colB(n + 1, 0);
        for(int i = 1; i <= n; i ++) {
            if(distA[i] == r && distB[i] == r + 1) colA[col[i]] ++;
            if(distA[i] == r + 1 && distB[i] == r) colB[col[i]] ++;
        }
        ll ans = 0;
        for(int i = 1; i <= n; i ++) ans += (ll)colA[i] * colB[i];
        cout << ans << "\n";
    }
}

int main() {
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int _ = 1; //cin >> _;
    while(_--) solve();
    return 0;
}