CF1681F Unique Occurrences

CF1681F Unique Occurrences

线段树分治练手题，半小时秒了。

考虑对每个颜色 w 计算有多少路径经过恰好一条颜色为 w 的边。

具体地，将所有颜色为 w 的边从原树中删去，对于每一条颜色为 w 的边的两端点连通块大小乘积的和即为这个颜色的答案。

将颜色看作时间线，那么一条颜色为 w 的边建在时间 [1,w-1] 和 [w+1,n] 上，线段树优化即可。

计算连通块大小用并查集即可，线段树分治考虑可撤销并查集。

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

#include <bits/stdc++.h>

using namespace std;

namespace Fread
{
    const int SIZE = 1 << 23;
    char buf[SIZE], *S, *T;
    inline char getchar()
    {
        if (S == T)
        {
            T = (S = buf) + fread(buf, 1, SIZE, stdin);
            if (S == T)
                return '\n';
        }
        return *S++;
    }
}

namespace Fwrite
{
    const int SIZE = 1 << 23;
    char buf[SIZE], *S = buf, *T = buf + SIZE;
    inline void flush()
    {
        fwrite(buf, 1, S - buf, stdout);
        S = buf;
    }
    inline void putchar(char c)
    {
        *S++ = c;
        if (S == T)
            flush();
    }
    struct NTR
    {
        ~NTR()
        {
            flush();
        }
    } ztr;
}

#ifdef ONLINE_JUDGE
#define getchar Fread::getchar
#define putchar Fwrite::putchar
#endif

//#define int long long
typedef long long ll;
#define ha putchar(' ')
#define he putchar('\n')

inline int read()
{
    int x = 0, f = 1;
    char c = getchar();
    while (c < '0' || c > '9')
    {
        if (c == '-')
            f = -1;
        c = getchar();
    }
    while (c >= '0' && c <= '9')
        x = x * 10 + c - '0', c = getchar();
    return x * f;
}

inline void write(ll x)
{
    if (x < 0)
    {
        putchar('-');
        x = -x;
    }
    if (x > 9)
        write(x / 10);
    putchar(x % 10 + '0');
}

const int _ = 5e5 + 10;

int n, x, y, z, sz[_], fa[_], top;

ll ans;

vector<pair<int, int>> d[_], e[_ << 2];

int find(int x)
{
    while (x ^ fa[x])
        x = fa[x];
    return x;
}

void upd(int o, int l, int r, int L, int R, pair<int, int> pi)
{
    if (L <= l && r <= R)
        return e[o].push_back(pi), void();
    int mid = (l + r) >> 1;
    if (L <= mid)
        upd(o << 1, l, mid, L, R, pi);
    if (R > mid)
        upd(o << 1 | 1, mid + 1, r, L, R, pi);
}

void solve(int o, int l, int r)
{
    vector<int> dl;
    for (auto v : e[o])
    {
        int p = find(v.first), q = find(v.second);
        if (p == q)
            continue;
        if (sz[p] > sz[q])
            swap(p, q);
        dl.push_back(p), sz[q] += sz[p], fa[p] = q;
    }
    if (l == r)
        for (auto v : d[l])
            ans += 1ll * sz[find(v.first)] * sz[find(v.second)];
    else
    {
        int mid = (l + r) >> 1;
        solve(o << 1, l, mid), solve(o << 1 | 1, mid + 1, r);
    }
    reverse(dl.begin(), dl.end());
    for (auto v : dl)
        sz[fa[v]] -= sz[v], fa[v] = v;
}

signed main()
{
    n = read();
    for (int i = 1; i < n; ++i)
    {
        x = read(), y = read(), z = read();
        d[z].push_back({x, y});
        if (z > 1)
            upd(1, 1, n, 1, z - 1, {x, y});
        if (z < n)
            upd(1, 1, n, z + 1, n, {x, y});
    }
    for (int i = 1; i <= n; ++i)
        sz[fa[i] = i] = 1;
    solve(1, 1, n);
    write(ans), he;
    return 0;
}