Atcoder Code Festvial 2017 Final J Tree MST

Problem

Atcoder

給你一棵 $n$ 個節點的樹，每個點有權值 $w[i]$ ，邊帶權。現構建一張完全圖，對於任意一對點 $(x, y)$ ，有一條長度爲 $w[x] + w[y]+ dis(x, y)$ 的邊。求這張圖的最小生成樹。

話說Code Festival在官網上咋進啊？qwqq

Solution

隊長：直覺告訴我這題是Boruvka算法
我：被神仙吊打.jpg

考慮Boruvka算法，我們就需要優化找到每個點最優選擇的點的過程。不妨考慮點分治，設 $sum[x]$ 表示 $x$ 到根的路徑權值和，那麼我們找到 $w[x]+sum[x]$ 最小的點，那麼如果mst中有一條邊跨越了當前根，這條邊的一個端點一定是 $x$ ，把子樹中的點都向 $x$ 連邊即可。因此我們就把可能的邊減少到了 $O(n\log n)$ 級別。

等等，同一棵子樹內的點對可能有問題？爲了處理方便，我們可以放縮條件，把邊權設爲兩點點權加上兩點到根的權值，這樣在同一棵子樹中的邊肯定不是最優的，我們只需要把這些可能的邊最後做一次kruskal即可。

正確性：點分治第一層的連邊就可以保證一定聯通，而我們又把可能的邊都存了起來，那麼mst顯然就已經被包括在這些邊之中了。

時間複雜度$O(n\log^2 n)$，空間複雜度$O(n\log n)$。

---

官方題解是這樣的：直接用Boruvka算法，對於一個點，我們想要找到連向另一個聯通塊的邊權最小的邊。dp記錄最優的邊權和相應的聯通塊編號，dp兩次即可得到每個點的最優選擇。而每次聯通塊個數至少減半，因此算法的總複雜度爲 $O(n\log n)$ 。

Code

#include <algorithm>
#include <cstdio>
#include <vector>
using namespace std;
typedef long long ll;
const int maxn=200010,INF=0x3f3f3f3f;
template <typename Tp> inline int getmin(Tp &x,Tp y){return y<x?x=y,1:0;}
template <typename Tp> inline int getmax(Tp &x,Tp y){return y>x?x=y,1:0;}
template <typename Tp> inline void read(Tp &x)
{
    x=0;int f=0;char ch=getchar();
    while(ch!='-'&&(ch<'0'||ch>'9')) ch=getchar();
    if(ch=='-') f=1,ch=getchar();
    while(ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    if(f) x=-x;
}
struct Edge{
	int u,v;ll w;
	Edge(const int _u=0,const int _v=0,const ll _w=0ll){u=_u;v=_v;w=_w;}
	bool operator < (const Edge &b)const{return w<b.w;}
};
struct data{int v,w,nxt;}edge[maxn<<1];
int n,p,sn,rt,top,a[maxn],head[maxn],sz[maxn],mx[maxn],vis[maxn],stk[maxn];
ll ans,sum[maxn];
vector<Edge> G;
vector<Edge>::iterator itr;
struct Dsu{
	int f[maxn];
	void init(int n){for(int i=1;i<=n;i++) f[i]=i;}
	int find(int x){return f[x]==x?x:f[x]=find(f[x]);}
	int merge(int x,int y)
	{
		int fx=find(x),fy=find(y);
		if(fx==fy) return 0;
		return f[fy]=fx;
	}
}dsu;
void insert(int u,int v,int w)
{
	edge[++p]=(data){v,w,head[u]};head[u]=p;
	edge[++p]=(data){u,w,head[v]};head[v]=p;
}
void getrt(int x,int pre)
{
	sz[x]=1;mx[x]=0;
	for(int i=head[x];i;i=edge[i].nxt)
	  if(!vis[edge[i].v]&&edge[i].v^pre)
	  {
	  	getrt(edge[i].v,x);
	  	sz[x]+=sz[edge[i].v];
	  	getmax(mx[x],sz[edge[i].v]);
	  }
	getmax(mx[x],sn-sz[x]);
	if(mx[x]<mx[rt]) rt=x;
}
void input()
{
	int u,v,w;
	read(n);
	for(int i=1;i<=n;i++) read(a[i]);
	for(int i=1;i<n;i++){read(u);read(v);read(w);insert(u,v,w);}
	dsu.init(n);
	sn=n;mx[rt=0]=INF;getrt(1,1);
}
void getdis(int x,int pre)
{
	stk[++top]=x;sz[x]=1;
	for(int i=head[x];i;i=edge[i].nxt)
	  if(!vis[edge[i].v]&&edge[i].v^pre)
	  {
	  	sum[edge[i].v]=sum[x]+edge[i].w;
	  	getdis(edge[i].v,x);
	  	sz[x]+=sz[edge[i].v];
	  }
}
void dfs(int x)
{
	int id;ll mn=1e18;
	vis[x]=1;top=0;sum[x]=0ll;
	getdis(x,x);
	for(int i=1;i<=top;i++)
	  if(getmin(mn,a[stk[i]]+sum[stk[i]]))
	    id=stk[i];
	for(int i=1;i<=top;i++)
	  G.push_back(Edge(id,stk[i],a[stk[i]]+sum[stk[i]]+a[id]+sum[id]));
	for(int i=head[x];i;i=edge[i].nxt)
	  if(!vis[edge[i].v])
	  {
	  	sn=sz[edge[i].v];rt=0;
	  	getrt(edge[i].v,x);
	  	dfs(rt);
	  }
}
int main()
{
	#ifndef ONLINE_JUDGE
	freopen("in.txt","r",stdin);
	#endif
	input();
	dfs(rt);
	sort(G.begin(),G.end());
	for(itr=G.begin();itr!=G.end();++itr)
	  if(dsu.merge(itr->u,itr->v))
	    ans+=itr->w;
	printf("%lld\n",ans);
	return 0;
}