【题解】CF1681F-Unique Occurrences

· · 题解

原文见我的博客。

给定一个 n 个结点的树，每条边有一个颜色，记 f(x,y) 表示结点 x 到 y 的路径上只出现了一次的颜色的数量，求 \sum\limits_{x < y}f(x, y) 。数据保证 n \le 5 \times 10^5 。

题面很短，解法很多。

由于我们是求所有 f(x,y) 的总和，根据经验我们将问题转化为对每个颜色 w 计算有多少路径经过恰好一条颜色为 w 的边。这样每种颜色是独立的，可以分开计算。

方法一：虚树

很自然我们可以想到虚树的解法，我们将颜色相同的边的两个端点保留下来，建立虚树后直接跑 DP。

方法二：二次扫描换根

如果我们设 DP 的状态 f_i 表示以 i 为根子树中，以 i 为起点，不经过颜色 w 的单链数量，同理 g_i 表示删去 i 的子树的数量，那么一条颜色为 w 的边的贡献就是两端 f,g 的乘积。f 可以在 DFS 的同时维护栈求得，然后二次扫描换根求得 g。

方法三：线段树分治

如果现在考虑颜色 w，那么删去所有该颜色的边，求得剩下每个连通块的大小，那么每条颜色 w 的边的贡献就是两端连通块大小的乘积。显然这个过程可以通过线段树分治快速维护。

方法四：线段树

固定一个根，如果我们把颜色为 w 的边的深度较深的点记为关键点，考虑用 DS 快速求出方法三中的连通块大小。

我们先将树上每个点赋值为 1，将关键点按 DFS 序排序，每次求得子树和记为连通块大小，然后将子树内所有点赋值为 0，这些都可以用数据结构快速维护。

代码实现的是方法三，因为好想好写，时间复杂度 \mathcal{O}(n\log ^2 n)。

#define N 500005
int n, fa[N], sz[N];
inline int get(int x){while(x != fa[x])x = fa[x]; return x;}
vector<Pr>u[N << 2], c[N];
#define ls (x << 1)
#define rs (ls | 1)
void ins(int x,int L,int R,int l,int r,Pr w){
    if(L >= l && R <= r){u[x].pb(w); return;}
    int mid = (L + R) >> 1;
    if(mid >= l)ins(ls, L, mid, l, r, w);
    if(mid < r)ins(rs, mid + 1, R, l, r, w);
}
LL ans;
void solve(int x,int l,int r){
    vector<int>dl;
    go(w, u[x]){
        int p = get(w.fi), q = get(w.se);
        if(sz[p] > sz[q])swap(p, q);
        dl.pb(p), sz[q] += sz[p], fa[p] = q;
    }
    if(l == r)go(w, c[l])ans += sz[get(w.fi)] * 1LL * sz[get(w.se)];
    else {
        int mid = (l + r) >> 1;
        solve(ls, l, mid), solve(rs, mid + 1, r);
    }
    reverse(dl.begin(), dl.end());
    go(w, dl)sz[fa[w]] -= sz[w], fa[w] = w;
}
int main() {
    read(n);
    rp(i, n - 1){
        int x, y, z; read(x, y, z);
        c[z].pb(mp(x, y));
        if(z > 1)ins(1, 1, n, 1, z - 1, {x, y});
        if(z < n)ins(1, 1, n, z + 1, n, {x, y});
    }
    rp(i, n)sz[fa[i] = i] = 1;
    solve(1, 1, n);
    printf("%lld\n", ans);
    return 0;
}