【題解】Codeforces 1213G. Path Queries

#include <iostream>
#include <vector>
#include <algorithm>
#include <cstring>
using namespace std;
 
int n, m, a, b, w, ord[200005], p[200005];
long long cnt, q[200005];
vector <pair<int, int> > v[200005];
 
int f(int x){
    if (p[x] < 0) return x;
    return p[x] = f(p[x]);
}
bool cmp(int x, int y){
    return q[x] < q[y];
}
 
int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n >> m;
    for (int i = 1; i < n; i++){
        cin >> a >> b >> w;
        v[w].push_back({a, b});
        p[i] = -1;
    }
    p[n] = -1;
    for (int i = 0; i < m; i++){
        cin >> q[i];
        ord[i] = i;
    }
    sort(ord, ord+m, cmp);
    int pre = 0;
    for (int i = 0; i < m; i++){
        for (int j = pre+1; j <= q[ord[i]]; j++){
            for (auto k:v[j]){
                int a = f(k.first), b = f(k.second);
                if (a != b){
                    if (-p[a] < -p[b]) swap(a, b);
                    cnt += (long long)(-p[a])*(-p[b]);
                    p[a] += p[b];
                    p[b] = a;
                }
            }
        }
        pre = q[ord[i]];
        q[ord[i]] = cnt;
    }
    for (int i = 0; i < m; i++){
        cout << q[i] << " ";
    }
}