1681F - Unique Occurrences - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 5e5 + 10;
int n,m,k,T;
#define endl '\n'

struct DSU
{
    int p[N];
    int sz[N];
    vector<pair<int &,int>> his_sz;
    vector<pair<int &,int>> his_fa;
    void init(int n)
    {
        for(int i = 1;i <= n;i++) p[i] = i,sz[i] = 1;
    }
    int find(int x)
    {
        while(x != p[x]) x = p[x];
        return x;
    }
    bool same(int u,int v)
    {
        return find(u) == find(v);
    }
    void merge(int u,int v)
    {
        int x = find(u),y = find(v);
        if(x == y) return ;
        if(sz[x] < sz[y]) std::swap(x,y);
        his_sz.push_back({sz[x],sz[x]});
        sz[x] += sz[y];
        his_fa.push_back({p[y],p[y]});
        p[y] = x;
    }
    void roll(int h)
    {
        while(his_fa.size() > h)
        {
            his_fa.back().first = his_fa.back().second;
            his_fa.pop_back();
            his_sz.back().first = his_sz.back().second;
            his_sz.pop_back();
        }
    }
}dsu;
vector<array<int,2>> g[N];

int dfs(int l,int r)
{
    if(l == r)
    {
        int res = 0;
        for(auto item : g[l])
        {
            res += dsu.sz[dsu.find(item[0])] * dsu.sz[dsu.find(item[1])];
        }
        return res;
    }
    int o = dsu.his_sz.size();
    int mid = (l + r) >> 1;
    int res = 0;
    for(int i = mid + 1;i <= r;i++)
    {
        for(auto item : g[i])
        {
            dsu.merge(item[0],item[1]);
        }
    }
    res += dfs(l,mid);
    dsu.roll(o);
    o = dsu.his_sz.size();
    for(int i = l;i <= mid;i++)
    {
        for(auto item : g[i])
        {
            dsu.merge(item[0],item[1]);
        }
    }
    res += dfs(mid + 1,r);
    dsu.roll(o);
    return res;
}

void solve()
{
    cin >> n;
    dsu.init(n);
    for(int i = 1;i < n;i++)
    {
        int a,b,w;
        cin >> a >> b >> w;
        g[w].push_back({a,b});
    }
    cout << dfs(1,n) << endl;
}

signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    T = 1;
    while(T--)
    {
        solve();
    }
}