kmjp's blog

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

Codeforces #591 Div1 C. Paint the Tree

この回なんかAC数の減少がなだらかだよね。
https://codeforces.com/contest/1240/problem/C

問題


N要素の木を成す無向グラフが与えられる。
各辺にはスコアが設定されている。

ここで、各頂点にK色ずつ割り当てることにする。
ある色は2頂点までしか振れない。一方色の種類は無限にある。
辺の両端が1色でも共有する場合、その辺のスコアを獲得できる。
最適な色の割り当てをしたとき、スコアの総和の最大値を求めよ。

解法

各頂点に対し、つながるK本以下の辺をスコアの加算対象とすることができることになる。
そこでDFSを行い、

  • 親方向の辺を含まずほかに子方向の辺をK本以下スコアに加算する場合
  • 親方向の辺を含み、ほかに子方向の辺を(K-1)本以下スコアに加算する場合

のそれぞれの最大スコアを木DPで求めていけばよい。

int Q;
int N,K;
vector<pair<int,int>> E[505050];

pair<ll,ll> dfs(int cur,int pre) {
	int pc=0;
	vector<pair<ll,ll>> V;
	
	FORR(e,E[cur]) if(e.first!=pre) {
		auto p=dfs(e.first,cur);
		p.first=max(p.first,p.second);
		p.second=max(p.first,p.second+e.second);
		swap(p.first,p.second);
		p.first-=p.second;
		V.push_back(p);
	}
	
	sort(ALL(V));
	reverse(ALL(V));
	pair<ll,ll> ret={0,0};
	int i;
	FOR(i,V.size()) {
		auto& v=V[i];
		if(i<K) ret.first+=v.first+v.second;
		else ret.first+=v.second;
		if(i<K-1) ret.second+=v.first+v.second;
		else ret.second+=v.second;
	}
	return ret;
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	scanf("%d",&Q);
	while(Q--) {
		scanf("%d%d",&N,&K);
		FOR(i,N) E[i].clear();
		
		FOR(i,N-1) {
			scanf("%d%d%d",&x,&y,&r);
			x--;
			y--;
			E[x].push_back({y,r});
			E[y].push_back({x,r});
		}
		
		auto ret=dfs(0,-1);
		cout<<ret.first<<endl;
		
		
	}
}

まとめ

わざわざ色がどうこうなんて言わなくてもよかったのになぁとも思う。