蜜汁TLE求调
查看原帖
蜜汁TLE求调
503792
Svemit楼主2024/7/17 18:43

rt,不知道是写挂了还是写丑了。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 2e5 + 5;

int n;
vector<pair<int, int>> e[N];
int siz[N], maxsz[N];
int a[N];
bool st[N];
pair<ll, int> q[N];
int qtot;
int fa[N];

int find(int x) {
	return fa[x] == x ? x : fa[x] = find(fa[x]);
}

struct Edge {
	int u, v;
	ll w;
	bool operator <(const Edge &t) const {
		return w < t.w;
	}
};
vector<Edge> E;

void get_rt(int u, int fa, int tot, int &rt) {
	siz[u] = 1;
	for (auto [v, w] : e[u]) {
		if (v == fa || st[v]) {
			continue;
		}
		get_rt(v, u, tot, rt);
		siz[u] += siz[v];
		maxsz[u] = max(maxsz[u], siz[v]);
	}
	maxsz[u] = max(maxsz[u], tot - siz[u]);
	if (maxsz[u] < maxsz[rt]) {
		rt = u;
	}
}

void get_siz(int u, int fa) {
	siz[u] = 1;
	for (auto [v, w] : e[u]) {
		if (v == fa || st[v]) {
			continue;
		}
		get_siz(v, u);
		siz[u] += siz[v];
	}
}

void get_dist(int u, int fa, ll d) {
	q[++ qtot] = {d + a[u], u};
	for (auto [v, w] : e[u]) {
		if (v == fa || st[v]) {
			continue;
		}
		get_dist(v, u, d + w);
	}
}

void solve(int u, int tot) {
	int rt = 0;
	get_rt(u, 0, tot, rt);
	u = rt;
	st[u] = true;
	// cout << "----------------\n";
	// cout << u << "\n";
	get_dist(u, 0, 0);
	sort(q + 1, q + 1 + qtot);	
	// for (int i = 1; i <= qtot; i ++) {
	// 	cout << q[i].first << ' ' << q[i].second << "\n";
	// }
	for (int i = 2; i <= qtot; i ++) {
		E.push_back({q[1].second, q[i].second, q[1].first + q[i].first});
	}
	qtot = 0;
	get_siz(u, 0);
	for (auto [v, w] : e[u]) {
		if (st[v]) {
			continue;
		}
		solve(v, siz[v]);
	}
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);

    maxsz[0] = 1e9;

    cin >> n;

    for (int i = 1; i <= n; i ++) {
    	fa[i] = i;
    }

    for (int i = 1; i <= n; i ++) {
    	cin >> a[i];
    }

    for (int i = 1; i < n; i ++) {
    	int u, v, w;
    	cin >> u >> v >> w;
    	e[u].emplace_back(v, w);
    	e[v].emplace_back(u, w);
    }

    solve(1, n);

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

    ll ans = 0;

    for (auto [u, v, w] : E) {
    	// cout << u << ' ' << v << ' ' << w << "\n";
    	u = find(u), v = find(v);
    	if (u == v) {
    		continue;
    	}
    	fa[u] = v;
    	ans += w;
    }

    cout << ans << "\n";

    return 0;
}
2024/7/17 18:43
加载中...