Codeforces 87 D. Beautiful Road

問題

重み付き無向木が与えられる。
このグラフの二点を結ぶパス全てについて、
パス上で重みが最も重い辺(複数ある場合全て)に木を一本植えるという操作をする。


全ての操作の後、最も木が植えられている辺に植えられている木の本数および、
その本数の木が植えられている辺の番号を全て出力せよ。

制約条件

n≦100000
重み≦10^9

方針

まずは、辺の重みが全てユニークな場合を考える。
これは、辺を重みの小さい順にソートして、
Union-Findでマージしながら、
辺の左右に属する頂点をそれぞれa,bとすると、その辺に植えられた木の本数は2*a*bと計算していくことができる。


辺の重みに同一のものがある場合、これだとうまく行かない。
ある木のある頂点について、右側でその頂点から(同じ重みの辺を使って)たどれる頂点と、左側でその頂点からたどれる頂点の数を効率的に求めたい。


これは、DPを少し工夫することでできる。
辺を2倍に増やして、(2向き付け通りの向き付けをする)
「その辺から先に属する頂点の数」をメモ化して探索してやればよい。


DPテーブルも工夫することで、使いまわすことができる。
dpiにいくつのiのときのDPでメモ化されたかを覚えておくことで、
初期化のコストがはぶける。

ソースコード

struct edge{
	int u,v,w,id;
	bool operator<(const edge &r)const{
		return w<r.w;
	}
};
edge es[100000];
int n;
int p[100000],sz[100000];
int root(int x){
	if(x==p[x])return x;
	return p[x]=root(p[x]);
}
void merge(int a,int b){
	a=root(a); b=root(b);
	if(a==b)return;
	p[a]=b;
	sz[b]+=sz[a];
}
int cnt,cur;
ll ans[100000],mx;
int dp[200000],dpi[200000];
vector<edge> graph[100000];

int dfs(int u,int prev,int id){
	if(dpi[id]==cur)return dp[id];
	
	int res=sz[root(u)];
	rep(i,graph[u].size()){
		int v=graph[u][i].v, nid=graph[u][i].id;
		if(v!=prev)res+=dfs(v,u,nid);
	}
	dpi[id]=cur;
	return dp[id]=res;
}
void run()
{
	cin>>n;
	rep(i,n-1){
		int u,v,w; cin>>u>>v>>w;
		es[i]=(edge){u-1,v-1,w,i};
		dpi[i]=dpi[i+n]=-1;
	}
	sort(es,es+n-1);
	rep(i,n)p[i]=i,sz[i]=1;
	mx=-1;
	
	rep(i,n-1){
		vector<edge> e;
		for(int j=i;j<n-1&&es[j].w==es[i].w;j++){
			e.pb(es[j]);
			int u=root(es[j].u), v=root(es[j].v), id=es[j].id;
			graph[u].pb((edge){ u,v,0,id });
			graph[v].pb((edge){ v,u,0,id+n });
		}
		cur=i;
		rep(j,e.size()){
			int u=root(e[j].u), v=root(e[j].v), id=e[j].id;
			ans[e[j].id]=2ll*dfs(v,u,id)*dfs(u,v,id+n);
			if(ans[e[j].id]>mx)mx=ans[e[j].id], cnt=0;
			if(ans[e[j].id]==mx)cnt++;
		}
		rep(j,e.size()){
			int u=root(e[j].u), v=root(e[j].v);
			graph[u].clear();
			graph[v].clear();
			merge(u,v);
		}
		i+=e.size()-1;
	}
	cout<<mx<<" "<<cnt<<endl;
	rep(i,n-1)if(ans[i]==mx)cout<<i+1<<" ";
	cout<<endl;
}