Path Queries

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

const ll nmax = 1e9+7;
const ll nmax2 = 998244353;

class DSU {
public:
    std::vector<ll> f, siz;
    DSU(ll n) : f(n + 1), siz(n + 1, 1) {
        std::iota(f.begin() + 1, f.end(), 1);
    }
    ll leader(ll x) {
        while (x != f[x]) x = f[x] = f[f[x]];
        return x;
    }
    bool same(ll x, ll y) {
        return leader(x) == leader(y);
    }
    void merge(ll x, ll y) {
        x = leader(x);
        y = leader(y);
        if (x == y) return;
        siz[x] += siz[y];
        f[y] = x;
    }
    ll size(ll x) {
        return siz[leader(x)];
    }
};

int main(){
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(0);
    std::cout.tie(0);

    ll n, m;
    std::cin >> n >> m;

    std::map<ll, std::vector<std::pair<ll, ll>>> mp;
    std::vector<ll> adj[n + 1];
    std::set<std::pair<ll, ll>> l;
    for (ll i = 1, x, y, z; i < n; i++){
        std::cin >> x >> y >> z;
        adj[x].push_back(y);
        adj[y].push_back(x);
        mp[z].push_back({x, y});
        l.insert({z, 0});
    }

    std::map<ll, ll> ans;
    ll q[m + 1];
    for (ll i = 1; i <= m; i++){
        std::cin >> q[i];
        l.insert({q[i], 1});

    }

    std::sort(l.begin(), l.end());
    DSU dsu(n);

    ll res = 0;
    for (auto i : l){
        if (i.second == 1){
            ans[i.first] = res;
        } else {
            for (auto j : mp[i.first]){
                ll A = j.first;
                ll B = j.second;
                ll siz_A = dsu.siz[A];
                ll siz_B = dsu.siz[B];
                res = res - (siz_A * (siz_A - 1)) / 2 - (siz_B * (siz_B - 1)) / 2 + ((siz_A + siz_B) * (siz_A + siz_B - 1)) / 2;
                dsu.merge(A, B);
            }
        }
    }

    for (ll i = 1; i <= m; i++){
        std::cout << ans[q[i]] << ' ';
    }


}