题解：P12849 [蓝桥杯 2025 国 A] 公路【数据强度待检验】

· · 题解

抽象扭曲魔怔做法。

对每种颜色分别计算答案，考虑断开所有该种颜色的边，那么共有 n(n-1)-\sum\limits_{S}|S|(|S|-1) 条路径上有该颜色，其中 S 为剩下的每个连通块。

于是现在问题变成了颜色 i 有生效区间 [1,i-1]\cup[i+1,n]，维护每个连通块的 |S|(|S|-1) 之和，使用线段树分治+并查集即可，复杂度为 \mathcal O(n\log^2 n)。

擦着过的。

#include <bits/stdc++.h>
using namespace std;
#define il inline
#define IOSIZE (1 << 20)
char buf[IOSIZE], *p1 = buf, *p2 = buf;
#define gc() (p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, IOSIZE, stdin), p1 == p2) ? EOF : *p1++)
il int read() { int x = 0; char c = 0; while (c < '0' || c > '9') c = gc(); while (c >= '0' && c <= '9') x = x * 10 + c - '0', c = gc(); return x; }
typedef long long ll;
const int N = 5e5 + 5;
int n, u[N], v[N], fa[N], siz[N], top, stk[N]; ll C[N], ans, cur;
basic_string<int> t[N << 2];
il int get(int x) { return x == fa[x] ? x : get(fa[x]); }
il void merge(int x, int y) {
    if ((x = get(x)) == (y = get(y))) return;
    if (siz[x] > siz[y]) swap(x, y);
    cur -= C[siz[x]], cur -= C[siz[y]];
    fa[x] = y, siz[y] += siz[x], cur += C[siz[y]];
    stk[++top] = x;
}
il void upd(int p, int l, int r, int L, int R, int x) {
    if (L <= l && r <= R) return t[p] += x, void();
    int mid = l + r >> 1;
    if (L <= mid) upd(p << 1, l, mid, L, R, x);
    if (R > mid) upd(p << 1 | 1, mid + 1, r, L, R, x);
}
il void roll(int t) {
    while (top > t) {
        int x = stk[top--], y = fa[x];
        cur -= C[siz[y]], fa[x] = x, siz[y] -= siz[x];
        cur += C[siz[x]], cur += C[siz[y]];
    }
}
il void solve(int p, int l, int r) {
    int tot = top;
    for (int i = 0, s = t[p].size(); i < s; i++) merge(u[t[p][i]], v[t[p][i]]);
    if (l == r) ans -= cur;
    else {
        int mid = l + r >> 1;
        solve(p << 1, l, mid), solve(p << 1 | 1, mid + 1, r);
    }
    roll(tot);
}
int main() {
    n = read();
    for (int i = 1, w; i < n; i++) {
        u[i] = read(), v[i] = read(), w = read();
        if (w > 1) upd(1, 1, n, 1, w - 1, i);
        if (w < n) upd(1, 1, n, w + 1, n, i);
    }
    for (int i = 1; i <= n; i++) fa[i] = i, siz[i] = 1, C[i] = (ll)i * (i - 1);
    ans = C[n] * n, solve(1, 1, n), cout << ans;
}