洛谷 P2971 [USACO10HOL] Cow Politics G 题解

怎么没有树上启发式合并的题解呢？我来发一篇吧！

简化题意

给定一棵 \(n\) 个点的树，每个点属于 \(k\) 种颜色之一（每种颜色至少有 2 个点）。求每种颜色中，任意两点间的最大距离。

核心思想

树上两点 \(u,v\) 间的距离为 \(dep_u+dep_v-2×dep_{lca(u,v)}\)。同色节点中的最大距离一定对应其中两个节点。考虑枚举它们的 LCA。

对每个节点 \(u\)，若它作为 LCA，则考虑它的所有子树中的同色节点，计算不同子树中的同色节点间的最大距离并更新答案。

算法步骤

使用树上启发式合并，维护每种颜色的深度集合。

1. 预处理：计算每个点的深度、子树大小、重儿子，和以它为根子树的 DFS 序区间。
2. 启发式合并：
   • 处理轻子树，处理后从集合中删除其节点。
   • 处理重子树。
   • 对每个轻子树，计算其节点与已处理节点的同色最大距离并更新答案，再将其加入集合。
   • 将当前节点自身当作一棵轻子树处理。

数据结构

对于每种颜色都开一个动态开点平衡树 multiset 维护这种颜色的节点的深度，支持插入、删除和查询最大值。由于深度可能有重复，因此不能用 set。

复杂度分析

对于时间复杂度，每个节点都会被操作 \(O(\log n)\) 次，每次操作 \(O(\log n)\)。一共 \(n\) 个节点。因此总时间复杂度为 \(O(n \log^2 n)\)，可以通过本题数据。

对于空间复杂度，由于所有 multiset 中最多保存 \(n\) 个节点的深度，再加上其他长度为 \(n\) 的数组，总空间复杂度为 \(O(n)\)。

AC Code

#include <bits/stdc++.h>
#define rept(i,a,b) for(int i(a);i<=b;++i)
#define eb emplace_back
using namespace std;
constexpr int N=2e5+5,K=1e5+5;
multiset<int> s[K];
vector<int> g[N];
int id[N],dep[N],l[N],r[N],siz[N],ch[N],a[N],ans[K],tot;
void dfs(int u,int pre){
    siz[u]=1,l[u]=++tot,id[tot]=u;
    for(int v:g[u]){
        if(v==pre) continue;
        dep[v]=dep[u]+1;
        dfs(v,u);
        siz[u]+=siz[v];
        if(siz[v]>siz[ch[u]]) ch[u]=v;
    }
    r[u]=tot;
}
void add(int l,int r){  // 插入一棵子树
    rept(i,l,r){
        int u=id[i];
        s[a[u]].insert(dep[u]);
    }
}
void sub(int l,int r){  // 删除一棵子树
    rept(i,l,r){
        int u=id[i];
        // 坑点：如果用s[a[u]].erase(dep[u])则会删除所有等于dep[u]的元素，并不是只删一个
        s[a[u]].erase(s[a[u]].find(dep[u]));
    }
}
void calc(int l,int r,int lca){  // 更新答案
    rept(i,l,r){
        int u=id[i];
        if(s[a[u]].empty()) continue;
        int d=*s[a[u]].rbegin();
        ans[a[u]]=max(ans[a[u]],d+dep[u]-(dep[lca]<<1));
    }
}
void dsu(int u,int pre){
    for(int v:g[u]){  // 处理轻儿子
        if(v==pre||v==ch[u]) continue;
        dsu(v,u);
        sub(l[v],r[v]);
    }
    if(ch[u]) dsu(ch[u],u);  // 重儿子
    for(int v:g[u]){
        if(v==pre||v==ch[u]) continue;
        calc(l[v],r[v],u);
        add(l[v],r[v]);
    }
    // 处理这个节点本身
    calc(l[u],l[u],u);
    add(l[u],l[u]);
}
int main(){
    cin.tie(0)->sync_with_stdio(0);
    int n,k,t;
    cin>>n>>k;
    rept(i,1,n){
        cin>>a[i]>>t;
        if(t) g[i].eb(t),g[t].eb(i);
    }
    dfs(1,0);
    dsu(1,0);
    rept(i,1,k){
        cout<<ans[i]<<'\n';
    }
    return 0;
}