九州大学プログラミングコンテスト2018 G - Tapu & Tapi 2

本番の出来は酷かったけど程よい難易度の問題が揃ってて良い感じのコンテストだった
J問題はまだよくわかってない

問題 (言い換え)

n頂点の木が与えられる
最初x個の頂点は赤で、y個の頂点は青で塗られており、残りの頂点は色が付いていない
この色が付いていない頂点を全て赤か青で塗り、最後に赤の頂点と青の頂点を結ぶ辺を全て削除する
削除する辺の重みの総和の最小値を求めよ

2 <= n <= 5*10^5
1 <= (辺の重み) <= 10^9

解法

dp[v][c]: 頂点vを色cで塗ったときの、頂点vを根とする部分木での問題の答え
として適当な頂点からメモ化再帰すればok
遷移は
dp[v][赤] = Σ min{dp[u][赤], dp[u][青] + (辺u-vの重み)} (u: vの直接の子)
dp[v][青]も同様

実装

q[]が各頂点のデフォの色
0が無色、1と2が赤と青に対応してる

rec(v,p,qq): 頂点pから頂点vに下ってきて、頂点vを色qqで塗る場合の答え

#define int long long

int n, x, y, a, b, c, q[500500], dp[500500][3];
vector<pair<int,int> > e[500500];

int rec(int v, int p, int qq) {
	if (q[v] && q[v]!=qq) return 1ll<<60;
	if (dp[v][qq]!=(1ll<<60)) return dp[v][qq];
	dp[v][qq] = 0;
	each(i,e[v]) if (i.first!=p)
		dp[v][qq] += min(rec(i.first,v,1) + i.second*(qq!=1), rec(i.first,v,2) + i.second*(qq!=2));
	return dp[v][qq];
}

signed main() {
	cin >> n >> x >> y;
	rep(i,n-1) {
		cin >> a >> b >> c;
		a--, b--;
		e[a].push_back({b,c}), e[b].push_back({a,c});
	}
	rep(i,x) {
		cin >> c;
		q[c-1] = 1;
	}
	rep(i,y) {
		cin >> c;
		q[c-1] = 2;
	}
	rep(i,500500) rep(j,3) dp[i][j] = 1ll<<60;
	cout << min(rec(0,-1,1), rec(0,-1,2)) << endl;
}

https://beta.atcoder.jp/contests/qupc2018/tasks/qupc2018_g