Learning Algorithms

アルゴリズムの勉強メモ

QUPC G. Tapu & Tapi 2

G - Tapu & Tapi 2

最終的にたぷちゃんとたぴちゃんの両方が存在するような連結成分が存在しなければOKなのでそれ以外の3つの状態をすべて持つことにします。

$dp1[u] :=$ $u$ を含む連結成分にたぷちゃんもたぴちゃんも存在しないもので有効な状態にする最小コスト
$dp2[u] :=$ $u$ を含む連結成分にたぷちゃんのみが存在するもので有効な状態にする最小コスト
$dp3[u] :=$ $u$ を含む連結成分にたぴちゃんのみが存在するもので有効な状態にする最小コスト

この問題では遷移のための他の状態は何も必要なくて以下のように遷移できてこの問題は解けました。$v$ は $u$ の子ども、 $c$ は $u, v$ 間のコストです。

$dp1[u] = \sum min(dp1[v], min(dp2[v], dp3[v]) + c)$
$dp2[u] = \sum min(min(dp1[v], dp2[v]), dp3[v] + c)$
$dp3[u] = \sum min(min(dp1[v], dp3[v]), dp2[v] + c)$

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; (i) < (int) (n); (i) ++)
using namespace std;

static const long long INFL = 0x3f3f3f3f3f3f3f3fLL;

int main() {
	int n, x, y;
	scanf("%d%d%d", &n, &x, &y);
	vector<vector<pair<int, int>>> g(n);
	rep(i, n - 1) {
		int a, b, c;
		scanf("%d%d%d", &a, &b, &c);
		a --, b --;
		g[a].push_back({b, c});
		g[b].push_back({a, c});
	}
	vector<int> tg(n);
	rep(i, x) {
		int a;
		scanf("%d", &a);
		a --;
		tg[a] = 1;
	}
	rep(i, y) {
		int a;
		scanf("%d", &a);
		a --;
		tg[a] = 2;
	}
	vector<long long> dp1(n), dp2(n), dp3(n);
	function<void (int, int)> dfs = [&](int u, int prev) {
		if (tg[u] != 0) dp1[u] = INFL;
		if (tg[u] == 1) dp2[u] = INFL;
		if (tg[u] == 2) dp3[u] = INFL;
		for (auto e : g[u]) {
			int v = e.first;
			if (v == prev) continue;
			int c = e.second;
			dfs(v, u);
			if (dp1[u] != INFL) dp1[u] += min(dp1[v], min(dp2[v], dp3[v]) + c);
			if (dp2[u] != INFL) dp2[u] += min(min(dp1[v], dp2[v]), dp3[v] + c);
			if (dp3[u] != INFL) dp3[u] += min(min(dp1[v], dp3[v]), dp2[v] + c);
		}
	};
	dfs(0, -1);
	printf("%lld\n", min({dp1[0], dp2[0], dp3[0]}));
        return 0;
}