AT_cf17_final_j Tree MST 题解

一个有趣的结论是，一个图，选出若干个边集，满足他们的并为原图边集。则必有他们每个集合的 MST 的边的并，再求一次 MST 为原图的 MST。

考虑点分治，设当前重心为 $u$，设 $d_j = w_j + dist(u,j)$，则有对于两个点 $u,v$，其完全图上边权 $\leq d_u + d_v$。

每次分治时，把 $d$ 的最小值求出，将这个点和所有其他在分治树内的点连边，最后求 MST 即可。

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

#define int long long

constexpr int N = 2e5 + 5, M = N * 20;

int n;
vector<int> w;
vector<vector<pair<int, int>>> G;
vector<int> dis, sz;
bitset<N> del;

struct Edge
{
	int u, v, w;
	Edge(int _u, int _v, int _w): u(_u), v(_v), w(_w){}
	Edge(){}
};

vector<Edge> eg;
set<pair<int, int> > cv;
int tot, wc;

auto get_sz(int u, int f) -> void
{
	sz[u] = 0;
	if (del[u]) return;
	sz[u] = 1;
	for (auto &j : G[u])
	{
		if (j.first != f)
		{
			get_sz(j.first, u);
			sz[u] += sz[j.first];
		}
	}
}

auto get_wc(int u, int f) -> void
{
	if (del[u]) return;
	int maxn = tot - sz[u];
	for (auto &j : G[u])
	{
		if (j.first ^ f)
		{
			get_wc(j.first, u);
			maxn = max(maxn, sz[j.first]);
		}
	}
	if (maxn <= (tot >> 1)) wc = u; 
}

vector<int> total;

auto get_dist(int u, int f, int w) -> void
{
	dis[u] = 0;
	if (del[u]) return;
	total.emplace_back(u);
	dis[u] = dis[f] + w;
	for (auto &j : G[u])
	{
		if (j.first != f)
		{
			get_dist(j.first, u, j.second);
		}
	}
}

auto solve(int u) -> void
{
	if (del[u]) return;
	cv.clear();
	get_sz(u, 0);
	tot = sz[u];
	wc = u;
	get_wc(u, 0);
	u = wc;
	del[u] = 1;
	dis[u] = 0;
	for (auto &j : G[u])
	{
		if (del[j.first]) continue;
		total.clear();
		total.shrink_to_fit();
		get_dist(j.first, u, j.second);
		vector<pair<int, int> > djb;
		for (auto &k : total)
		{
			dis[k] += w[k];
			djb.emplace_back(make_pair(dis[k], k));
		}
		for (auto &k : djb)
		{
			cv.insert(make_pair(k.first, k.second));
		}
	}
	for (auto it = cv.begin(); it != cv.end(); ++it)
	{
		if (it != cv.begin())
		{
			eg.emplace_back(Edge((*it).second, (*(cv.begin())).second, (*it).first + (*(cv.begin())).first));
		}
		eg.emplace_back(u, (*(it)).second, (*it).first + w[u]);
	}
	for (auto &j : G[u]) solve(j.first);
}

class Union_Find
{
public:
	vector<int> fa;
	auto Init() -> void
	{
		fa.resize(n + 1);
		for (int i = 0; i <= n; i++) fa[i] = i;
	}
	auto find(int u) -> int
	{
		return (fa[u] == u ? u : fa[u] = find(fa[u]));
	}
	auto merge(int u, int v) -> void
	{
		fa.at(find(u)) = find(v);
	}
}uf;

auto main() -> signed
{
	ios::sync_with_stdio(0), cin.tie(0);
	cin >> n;
	del.reset();
	w.resize(n + 1);
	dis.resize(n + 1);
	sz.resize(n + 1);
	G.resize(n + 1);
	for (int i = 1; i <= n; i++) cin >> w[i];
	for (int i = 1; i < n; i++)
	{
		int u, v, w;
		cin >> u >> v >> w;
		G[u].emplace_back(make_pair(v, w));
		G[v].emplace_back(make_pair(u, w)); 
	}
	solve(1);
	uf.Init();
	long long ans = 0LL;
	sort(eg.begin(), eg.end(), [&](const Edge& x, const Edge& y){return x.w < y.w;});
	for (int i = 0; i < eg.size(); i++)
	{
		auto p = eg.operator[](i);
		if (uf.find(p.u) ^ uf.find(p.v))
		{
			ans += p.w;
			uf.merge(p.u, p.v);
		}
	}
	cout << ans << "\n";
	return 0;
}