kmjp's blog

競技プログラミング参加記です

Codeforces #543 Div1 D. Power Tree

問題設定がややこしい。
https://codeforces.com/contest/1120/problem/D

問題

根付き木が与えられる。
各頂点にはコストが設定されている。

ここで、いくつかの頂点を選ぶ。
その後、各頂点に対し任意の整数を設定する。
そうすると、その頂点のSubTree内の葉には、その整数が加算されるとする。

さて、各葉に任意の整数を設定できるような頂点の選び方のうち、総コストの最小値を求めよ。

解法

先にオイラーツアーを実施しておき、葉の頂点に訪問順に連番を振ろう。
その上で、各頂点vのSubTree内にある葉の連番区間[L(v),R(v)]を求めておく。
ここで、仮にvを選ぶということは、L(v)とR(v)+1の葉頂点の間に、任意の差をつけられるということに相当する。

さて、全ての葉に任意の値を設定できるようにするには、上記差をつけられる関係の頂点間に辺を張ったとき、それが全域木を成すことである。
よって上記辺が元の頂点コストを持つものとみなし、最小全域木を求めればよい。

int N;
int C[202020];
vector<int> E[202020];
map<int,vector<int>> M;
int id;
int L[202020],R[202020];

set<int> S,leafs;

template<int um> class UF {
	public:
	vector<int> par,rank;
	UF() {rank=vector<int>(um,0); for(int i=0;i<um;i++) par.push_back(i);}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<500000> uf;

void dfs(int cur,int pre) {
	L[cur]=id;
	if(cur>0 && E[cur].size()==1) {
		leafs.insert(L[cur]);
		id++;
	}
	
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur);
	R[cur]=id;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>C[i];
		M[C[i]].push_back(i);
	}
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,-1);
	leafs.insert(N);
	
	ll ret=0;
	FORR(m,M) {
		FORR(v,m.second) if(uf[L[v]]!=uf[R[v]]) S.insert(v);
		FORR(v,m.second) if(uf[L[v]]!=uf[R[v]]) {
			ret+=m.first;
			uf(L[v],R[v]);
		}
	}
	cout<<ret<<" "<<S.size()<<endl;
	FORR(s,S) cout<<(s+1)<<" ";
}

まとめ

区間に関する操作を、2点間の差とみなす、という部分に本番気付けなかったのが失敗。