G_Path_Queries.cpp

#include <bits/stdc++.h>
using namespace std;

class DSU {
public:
    vector<int> par, sz;
    DSU(int n){
        par.resize(n + 1);
        iota(par.begin(), par.end(), 0);
        sz.resize(n + 1, 1);
    }
    int find(int x){
        return x == par[x] ? x : par[x] = find(par[x]);
    }
    bool merge(int u, int v) {
        int p1 = find(u), p2 = find(v);
        if (p1 == p2) return 0;
        if (sz[p1] < sz[p2]) swap(p1, p2);
        par[p2] = p1;
        sz[p1] += sz[p2];
        return 1;
    }
};

int32_t main(){
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int n, q;
    cin >> n >> q;
    vector<array<int, 3>> edges;
    for(int i=0;i<n-1;i++) {
        int u, v, w;
        cin >> u >> v >> w;
        edges.push_back({u, v, w});
    }

    sort(edges.begin(), edges.end(), [&](const array<int, 3> &a, const array<int, 3> &b){
        return a[2] < b[2];
    });

    vector<pair<int, int>> qry;
    for(int i=0;i<q;i++) {
        int w;
        cin >> w;
        qry.push_back({w, i});
    }

    sort(qry.begin(), qry.end());

    DSU ds(n + 1);
    int j = 0;
    vector<long long> ans(q);
    long long cnt = 0;
    for(int i=0;i<q;i++) {
        int limit = qry[i].first;
        int id = qry[i].second;
        while(j < n - 1 && edges[j][2] <= limit) {
            int u = edges[j][0], v = edges[j][1];
            int t1 = ds.find(u);
            int t2 = ds.find(v);
            if(t1 != t2) {
                cnt += 1LL * ds.sz[t1] * ds.sz[t2];
                ds.merge(u, v);
            }
            j++;
        }
        ans[id] = cnt;
        
    }

    for(int i=0;i<q;i++) cout << ans[i] << " ";


    return 0;
}