洛谷P2971 [USACO10HOL] Cow Politics G 题解 树上启发式合并（dsu on tree）

题目链接：https://www.luogu.com.cn/problem/P2971

这里开了一个 map 数组：

map<int, int> mp[maxn];

其中：\(mp[u][x]\) 表示以 \(u\) 为根的子树中在政党 \(x\) 中的节点的最大深度。

由于我的处理方式并没有使用太多全局的数据处理数据，所以不需要做太多清楚的操作。

也就是说，轻儿子不需要搜两遍（因为数据都保存在 map 中了）。

这里的核心操作是：每个节点 \(u\) 和它的重儿子共用一个 map。 也就是说，一条重链上的点共用一个 map。

时间复杂度 \(O(n \log ^ 2 n)\)。

示例程序：

#include <bits/stdc++.h>
using namespace std;
const int maxn = 2e5 + 5;

int n, K, rt, a[maxn], id[maxn], sz[maxn], son[maxn], ans[maxn];
vector<int> g[maxn];
map<int, int> mp[maxn]; // mp[u][x] 表示以u为根的子树中在政党x中的节点的最大深度

void init() {
    for (int i = 1; i <= n; i++) {
        id[i] = i;
    }
}

void getsz(int u, int p) {
    sz[u] = 1;
    for (auto v : g[u]) {
        if (v != p) {
            getsz(v, u);
            sz[u] += sz[v];
            if (sz[v] > sz[ son[u] ])
                son[u] = v;
        }
    }
}

void dfs(int u, int p, int d) {
    for (auto v : g[u])
        if (v != p)
            dfs(v, u, d+1);
    if (son[u])
        id[u] = id[ son[u] ];
    int x = id[u];
    ans[ a[u] ] = max(ans[ a[u] ], mp[x][ a[u] ] - d);
    mp[x][ a[u] ] = max(mp[x][ a[u] ], d);
    for (auto v : g[u]) {
        if (v != p && v != son[u]) {
            for (auto &[a1, d1] : mp[ id[v] ]) {
                if (mp[x][a1])
                    ans[a1] = max(ans[a1], mp[x][a1] + d1 - 2 * d);
            }
            for (auto &[a1, d1] : mp[ id[v] ]) {
                mp[x][a1] = max(mp[x][a1], d1);
            }
            mp[ id[v] ].clear();
        }
    }

}

int main() {
    scanf("%d%d", &n, &K);
    for (int i = 1, p; i <= n; i++) {
        scanf("%d%d", a+i, &p);
        if (!p) rt = i;
        else g[p].push_back(i);
    }
    init();
    getsz(rt, -1);
    dfs(rt, -1, 0);
    for (int i = 1; i <= K; i++)
        printf("%d\n", ans[i]);
    return 0;
}