CF1681F Unique Occurrences

首先考虑拆贡献，\((u, v)\) 的边颜色为 \(w\)。主要思路就是考虑删去所有颜色为 \(w\) 的边，那么 \((u, v)\) 的贡献就是 \(u, v\) 所在连通块大小之积。

线段树分治

颜色为 \(w\) 的边在 \([1, w - 1] \cup [w + 1, n]\) 中出现，线段树分治维护。时间复杂度 \(\mathcal{O}(n \log^2 n)\)。

namespace Loop1st {
int n, fa[N], sz[N], stk[N], top;
ll ans;
vector<pii>c[N << 2], buc[N];
int find(int u) { return u == fa[u] ? u : find(fa[u]); }
void merge(int x, int y) {
    int fx = find(x), fy = find(y);
    if (sz[fx] < sz[fy]) swap(fx, fy);
    stk[++top] = fy;
    fa[fy] = fx;
    sz[fx] += sz[fy];
}
#define ls (u << 1)
#define rs (u << 1 | 1)
void add(int x, int y, pii e, int l = 1, int r = n, int u = 1) {
    if (x <= l && r <= y) { c[u].push_back(e); return ; }
    int mid = (l + r) >> 1;
    if (x <= mid) add(x, y, e, l, mid, ls);
    if (mid < y) add(x, y, e, mid + 1, r, rs);
}
void solve(int l = 1, int r = n, int u = 1) {
    int las = top;
    for (auto [x, y] : c[u]) {
        merge(x, y);
    }
    if (l == r) {
        for (auto [x, y] : buc[l]) {
            ans = ans + (ll)sz[find(x)] * sz[find(y)];
        }
    } else {
        int mid = (l + r) >> 1;
        solve(l, mid, ls); solve(mid + 1, r, rs);
    }
    while (top > las) {
        int x = stk[top--];
        sz[fa[x]] -= sz[x];
        fa[x] = x;
    }
}
void main() {
    cin >> n;
    for (int i = 1, u, v, w; i < n; i++) {
        cin >> u >> v >> w; buc[w].push_back({u, v});
    }
    for (int i = 1; i <= n; i++) sz[i] = 1, fa[i] = i;
    for (int i = 1; i <= n; i++) {
        for (auto e : buc[i]) {
            if (i > 1) add(1, i - 1, e);
            if (i < n) add(i + 1, n, e);
        }
    }
    solve();
    cout << ans << '\n';
}

}

虚树

把 \(\text{col}(x, y) = w\) 的 \(y\) 作为关键点，虚树上跑 DP。时间复杂度 \(\mathcal{O}(n \log n)\)。

线段树

将 \(\text{col}(x, y) = w\) 的点 \(y\) 按 \(\text{dfn}\) 排序，每次计算出当前点的子树中有多少 \(1\)，再将子树内所有点变为 \(0\)，线段树即可做到 \(\mathcal{O}(n \log n)\)。

换根 DP

用一个栈找到 \(u\) 子树内和 \((u, v)\) \(\text{dfn}\) 最近且颜色相同的边，就可以 DP 并统计答案了。时间复杂度 \(\mathcal{O}(n)\)。