#include <bits/stdc++.h>
using namespace std;
using ll = long long;
struct DSU {
int maxn;
vector<int> pa;
vector<int> sz;
DSU(int maxn = 0) { init(maxn); }
void init(int _maxn) {
maxn = _maxn;
if ((int)pa.size() < maxn + 5) {
pa.resize(maxn + 5);
sz.resize(maxn + 5);
}
reset(maxn);
}
// reset to fresh DSU of size n (1..n)
void reset(int n) {
// note: loop is necessary (same complexity as original) but avoids reallocations
for (int i = 0; i <= n; ++i) {
pa[i] = i;
sz[i] = 1;
}
}
int find(int x) {
while (x != pa[x]) {
pa[x] = pa[pa[x]];
x = pa[x];
}
return x;
}
bool merge(int x, int y) {
x = find(x);
y = find(y);
if (x == y) return false;
if (sz[x] < sz[y]) swap(x, y);
pa[y] = x;
sz[x] += sz[y];
return true;
}
bool same(int x, int y) { return find(x) == find(y); }
int size(int x) { return sz[find(x)]; }
};
struct Edge {
int u, v;
int w;
};
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, m;
if (!(cin >> n >> m)) return 0;
vector<Edge> e(m + 1);
for (int i = 1; i <= m; ++i) {
cin >> e[i].u >> e[i].v >> e[i].w;
}
sort(e.begin() + 1, e.end(), [](const Edge &a, const Edge &b) {
return a.w < b.w;
});
// preallocate lt/rt
vector<vector<int>> lt(m + 1), rt(m + 1);
DSU dsu(n); // allocate once
dsu.init(n);
// build lt (1..m), lt[0] stays empty
for (int i = 1; i <= m; ++i) {
dsu.reset(n);
dsu.merge(e[i].u, e[i].v);
lt[i].reserve((i > 0 ? lt[i - 1].size() : 0) + 1);
lt[i].push_back(i);
const auto &prev = lt[i - 1];
for (int it : prev) {
if (dsu.merge(e[it].u, e[it].v)) lt[i].push_back(it);
}
}
// build rt ; rt[m] stays empty in original logic (we mirror original indexing)
for (int i = m; i >= 1; --i) {
dsu.reset(n);
dsu.merge(e[i].u, e[i].v);
rt[i - 1].reserve((i <= m ? rt[i].size() : 0) + 1);
rt[i - 1].push_back(i);
const auto &next = rt[i];
for (int it : next) {
if (dsu.merge(e[it].u, e[it].v)) rt[i - 1].push_back(it);
}
}
int Q;
cin >> Q;
int p = 0; // number of edges with weight <= x (keeps same meaning as your original code)
for (int qi = 0; qi < Q; ++qi) {
int x;
cin >> x;
while (p + 1 <= m && e[p + 1].w <= x) ++p;
dsu.reset(n);
ll ans = 0;
const auto &L = lt[p];
const auto &R = rt[p];
int p1 = 0, p2 = 0;
int total = (int)L.size() + (int)R.size();
// merge edges in the order you had (closest weight to x first)
for (int t = 0; t < total; ++t) {
if (p2 == (int)R.size() || (p1 < (int)L.size() && ll(x - e[L[p1]].w) <= ll(e[R[p2]].w - x))) {
int idx = L[p1++];
if (dsu.merge(e[idx].u, e[idx].v)) ans += ll(x - e[idx].w);
} else {
int idx = R[p2++];
if (dsu.merge(e[idx].u, e[idx].v)) ans += ll(e[idx].w - x);
}
}
cout << ans << '\n';
}
return 0;
}
//chatgpt test
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |