Submission #964039

#TimeUsernameProblemLanguageResultExecution timeMemory
964039Gromp15Reconstruction Project (JOI22_reconstruction)C++17
42 / 100
5062 ms403784 KiB
#include <bits/stdc++.h>
#define ll long long
#define ar array
#define db double
#define all(x) x.begin(), x.end()
#define sz(x) (int)x.size()
using namespace std;
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
#define rint(l, r) uniform_int_distribution<int>(l, r)(rng)
template<typename T> bool ckmin(T &a, const T &b) { return a > b ? a = b, 1 : 0; }
template<typename T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }
struct dsu {
    vector<int> p, sze;
	int comps;
    dsu(int n) : p(n+1), sze(n+1), comps(n) {
        for (int i = 0; i <= n; i++) {
            p[i] = i, sze[i] = 1;
        }
    }
    int get(int v) {
        return v==p[v]?v:p[v]=get(p[v]);
    }
    int size(int v) {
        return sze[get(v)];
    }
    bool merge(int a, int b) {
        if ((a=get(a)) == (b=get(b))) return 0;
        if (sze[a]<sze[b]) swap(a, b);
        sze[a] += sze[b], p[b] = a;
		comps--;
        return 1;
    }
};
void test_case() {
	int n, m;
	cin >> n >> m;
	vector<ar<int, 3>> edges(m);
	for (auto &x : edges) cin >> x[0] >> x[1] >> x[2];
	vector<int> lst;
	vector<int> order(m);
	auto cmp = [&](int x, int y) {
		return edges[x][2] != edges[y][2] ? edges[x][2] < edges[y][2] : x < y;
	};
	iota(all(order), 0);
	sort(all(order), cmp);
	vector<vector<int>> suff(m), pref(m);
	{
		dsu d(n);
		vector<int> cur;
		for (int i = 0; i < m; i++) {
			auto [x, y, w] = edges[order[i]];
			if (d.merge(x, y)) {
				cur.emplace_back(order[i]);
			}
		}
		suff[0] = cur;
		for (int i = 1; i < m; i++) {
			auto mn = min_element(all(cur), cmp);
			if (*mn == order[i-1]) {
				d = dsu(n);
				cur.erase(mn);
				for (int idx : cur) {
					d.merge(edges[idx][0], edges[idx][1]);
				}
				for (int j = i; j < m; j++) {
					int pos = order[j];
					if (d.merge(edges[pos][0], edges[pos][1])) {
						cur.insert(upper_bound(all(cur), pos, cmp), pos);
						break;
					}
				}
			}
			suff[i] = cur;
		}
	}
	{
		dsu d(n);
		vector<int> cur;
		for (int i = m-1; i >= 0; i--) {
			auto [x, y, w] = edges[order[i]];
			if (d.merge(x, y)) {
				cur.emplace_back(order[i]);
			}
		}
		reverse(all(cur));
		assert(is_sorted(all(cur), cmp));
		pref[m-1] = cur;
		for (int i = m-2; i >= 0; i--) {
			auto mx = max_element(all(cur), cmp);
			if (*mx == order[i+1]) {
				d = dsu(n);
				cur.erase(mx);
				for (int idx : cur) {
					d.merge(edges[idx][0], edges[idx][1]);
				}
				for (int j = i; j >= 0; j--) {
					int pos = order[j];
					if (d.merge(edges[pos][0], edges[pos][1])) {
						cur.insert(upper_bound(all(cur), pos, cmp), pos);
						break;
					}
				}
			}
			pref[i] = cur;
		}
	}
	int q; cin >> q;
	const int INF = 2e9;
	while (q--) {
		int x; cin >> x;
		int l = 0, r = m-1, ans = -1;
		while (l <= r) {
			int mid = (l+r)/2;
			if (x >= edges[order[mid]][2]) ans = mid, l = mid+1;
			else r = mid-1;
		}
		int a = ans >= 0 ? sz(pref[ans]) - 1 : 0, b = 0, added = 0;
		dsu d(n);
		ll ans2 = 0;
		while (added < n-1) {
			int costA = ans >= 0 && a >= 0 ? x - edges[pref[ans][a]][2] : INF;
			int costB = ans + 1 < m && b < sz(suff[ans+1]) ? edges[suff[ans+1][b]][2] - x : INF;
			assert(costA >= 0);
			assert(costB >= 0);
			if (costA < costB) {
				auto [u, v, w] = edges[pref[ans][a]];
				if (d.merge(u, v)) added++, ans2 += costA;
				a--;
			}
			else {
				auto [u, v, w] = edges[suff[ans+1][b]];
				if (d.merge(u, v)) added++, ans2 += costB;
				b++;
			}
		}
		cout << ans2 << '\n';
	}
}
int main() {
    cin.tie(0)->sync_with_stdio(0);
    int t = 1;
//    cin >> t;
    while (t--) test_case();
}

#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...