CF1213G WA

#include <assert.h>
#include <bits/stdc++.h>
using namespace std;
#define dbg(...) logger(#__VA_ARGS__, __VA_ARGS__)
template <typename... Args> void logger(string vars, Args &&...values)
{
    cerr << vars << " = ";
    string delim = "";
    (..., (cerr << delim << values, delim = ", "));
    cerr << endl;
}

template <class T> using mpq = priority_queue<T, vector<T>, greater<T>>;
using ll = long long;
using pii = pair<int, int>;
using piii = pair<int, pii>;

int n, m;
piii es[200003];
pii q[200003];

int fa[200003], cnt[200003];
void init(int n)
{
    memset(fa, -1, sizeof(int) * (n + 1));
    for (int i = 0; i <= n + 1; ++i)
        cnt[i] = 1;
}
int find(int x)
{
    if (fa[x] == -1)
        return x;
    int px = find(fa[x]);
    cnt[px] += cnt[x];
    cnt[x] = 0;
    return fa[x] = px;
}
void join(int x, int y)
{
    int px = find(x), py = find(y);
    if (px == py)
        return;
    cnt[px] += cnt[py];
    cnt[py] = 0;
    fa[py] = px;
}

int main(int argc, char **argv)
{
    scanf("%d%d", &n, &m);
    init(n);
    mpq<piii> es_pq;
    for (int i = 1; i < n; ++i)
        scanf("%d%d%d", &es[i].second.first, &es[i].second.second, &es[i].first), es_pq.push(es[i]);

    for (int i = 1; i <= m; ++i)
        scanf("%d", &q[i].first), q[i].second = i;
    sort(q + 1, q + 1 + m);

    vector<ll> ans(m + 1);
    ll acc = 0;
    for (int i = 1; i <= m; ++i) {
        ll round = 0;
        set<int> gp;
        while (es_pq.empty() == false && es_pq.top().first <= q[i].first) {
            auto h = es_pq.top();
            es_pq.pop();
            auto [u, v] = h.second;
            join(u, v);
            gp.insert(u), gp.insert(v);
        }
        set<int> uniq_gp;
        for (auto n : gp)
            uniq_gp.insert(find(n));
        for (auto n : uniq_gp)
            round += ll(cnt[n]) * (cnt[n] - 1) / 2;
        acc = max(acc, round);
        ans[q[i].second] = acc;
    }
    for (int i = 1; i <= m; ++i)
        printf("%lld%c", ans[i], " \n"[i == m]);
    return 0;
};