【競プロ典型90問】073 - We Need Both a and b（★5）解説

問題文の要約

$N$頂点の木が与えられます。各頂点には$a$,$b$いずれかの文字が書かれています。$0$本以上の辺を削除した後、すべての連結成分が$a$,$b$両方の文字を含むものは何通りか求めて、$10^{9}+7$で割った余りを出力してください。

制約:$2\leq N\leq 10^{5}$

この解説を見るまでに必要な知識

そもそもこの問題がDPに見えるかの話

DP配列をどのように定義するか

初期値は何か

遷移式

頂点iにaが書かれている場合

頂点iにbが書かれている場合

実装

int n;
cin >> n;
vector<char> c(n);
cin >> c;
Graph g(n);
for (int i = 0; i < n - 1; ++i) {
	int a, b;
	cin >> a >> b;
	--a, --b;
	g.wadd(a, b);
}
//dp[pos][st] := 頂点posの連結成分について、それぞれの連結成分の状態がすべてstであるときの通り数
//st 0=aのみ,1=bのみ,2=aとb両方ある
//res = dp[0][2]
//mint は modint の略で、自動的にmodを取ってくれる整数型
vector<vector<mint>> dp(n, vector<mint>(3));

//頂点vから木dp
auto dfs = [&](auto&& self, int v, int prev) ->void {
	mint m1 = 1, m2 = 1;
    //頂点vのすべての子頂点sに対して計算
	for (auto [s, w] : g[v]) {
		if (s == prev)continue;
		self(self, s, v);
		if (c[v] == 'a') {
			m1 *= (dp[s][0] + dp[s][2]);
		}
		else {
			m1 *= (dp[s][1] + dp[s][2]);
		}
		m2 *= (dp[s][0] + dp[s][1] + 2 * dp[s][2]);
	}
	if (c[v] == 'a') {
		dp[v][0] = m1;
	}
	else {
		dp[v][1] = m1;
	}
	dp[v][2] = m2 - m1;
};

dfs(dfs, 0, -1);


cout << dp[0][2].val() << endl;

提出コード

終わりに