Magicode logo
Magicode
6 min read

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

https://cdn.apollon.ai/media/notebox/0358d657-b669-4277-a4f2-b458fa664564.jpeg

問題文の要約

NN頂点の木が与えられます。各頂点にはaa,bbいずれかの文字が書かれています。00本以上の辺を削除した後、すべての連結成分がaa,bb両方の文字を含むものは何通りか求めて、109+710^{9}+7で割った余りを出力してください。
制約:2N1052\leq N\leq 10^{5}
問題文へのリンク

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

 木DPについて、基本問題が解けるくらいだと嬉しいです、結局この記事で書きたいのは遷移式の話なので...

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

 この問題のような、通り数を求める問題は単純な全探索はもちろん間に合いません。よって、ある程度“状態をまとめて”数え上げる必要があります。この“状態をまとめる”というのはまさしく DP の得意とするところです。実は、数え上げ問題と DP は非常に相性が良いです。数え上げ問題を見たときはまず DP を考えてみるくらいの心構えで良いと思います。
 この問題で DP を考えた場合、グラフ構造が木なので、もちろん木 DP ですね。

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

 まず、DP で持たなければならない状態を考えてみます。木 DP を考えているので、必ず必要な“どの部分木を考えているか”というのはもちろん持たなければなりません。
 では、どの部分木かという状態だけで答えを求めるのに十分でしょうか。次のような例を考えてみましょう。次の2つの木はどちらも00通りです。しかし、この木が頂点aaの部分木であったときを考えます。部分木の情報から答えが何通りかを決定することを考えます。
 頂点全てにaaが書かれている部分木はどうやっても条件を満たさないので 00通りとなります。しかし、頂点全てにbbが書かれている部分木は、どの辺も切断しないことで条件を満たし、11通りとなります。同じ状態から別の更新結果が生まれてしまったので、持つ状態が不足していることが分かります。
 通り数は既に持っているので、追加で、“部分木に書かれているのがaaのみ、bbのみ、aabb両方含む”という情報を持っておく必要があることが分かります。
 これらの考察から、次のような DP を考えましょう。
dp[i][j] := 根が iの部分木の各連結成分について、状態がすべて j = {0:a のみ, 1:b のみ, 2:aとb両方含む}であるときの、辺を0本以上切断する方法の数
 ちなみに、状態が0または1(一種類の文字しか含まない)であるときの辺を切断する方法の数は、1通り(どれかの辺を切断するともう一方の文字を含む可能性がなくなる連結成分が生まれるため)しかないことが分かります。

初期値は何か

 木 DPなので葉の頂点 cの初期値を考えます。頂点 c に a が書かれているときは dp[c][0] = 1、頂点cにbが書かれているときはdp[c][1] = 1で、ほかの状態はすべて0となります。

遷移式

 頂点iの子頂点の情報を使ってdp[i][0],dp[i][1],dp[i][2]を更新していくことを考えましょう。

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

 dp[i][0] (頂点iの部分木に書かれている文字がすべてa)はどのように更新できるでしょうか。 頂点iの子頂点の一つをbjb_jとします。辺ibji-b_jを残すか切断するか考えましょう。
 もしbjb_jの状態が0(aのみが書かれている)である場合、その辺は必ず残さなければいけません。
 bjb_jの部分木の状態が1(bのみが書かれている)の場合、辺を削除しても残したままにしても条件を満たさないことがわかります。
 bjb_jの部分木の状態が2(aとb両方書かれている)であるとき、辺i-bjb_jは必ず切断しなければなりません。
 よって、dp[i][0]を更新するのには、dp[bjb_j][0]とdp[bjb_j][2]が必要であることが分かります。部分木の状態は一通りしかありえないので、この二つのうちどちらかは0です。片方しか要らないのですが、場合分けするのも面倒なので+を取るかmaxを取ってしまいましょう。
 次に、dp[i][1] (頂点iの部分木に書かれている文字がすべてb)を考えましょう。しかし、これはそもそも起こりえないことがわかります。頂点iにaが書かれているのに、根がiの部分木にbだけが書かれていることは起こりえません。よって、dp[i][1]は更新する必要はありません。
 最後に、dp[i][2] (頂点iの部分木にaとb両方含まれる)を考えましょう。もしbjb_jの状態が0(aのみが書かれている)である場合、その辺は必ず残さなければいけません。
bjb_jの部分木の状態が1(bのみが書かれている)の場合も、辺を残さなければいけないことがわかります。
bjb_jの部分木の状態が2(aとb両方書かれている)であるときはどうでしょうか。このとき、辺ibji-b_jは切断しても残したままにしても良いですね。 よって、+を取る前に2倍してやります。 以上から、dp[i][2]を更新するのに必要な部分木の情報は、dp[bjb_j][0]+dp[bjb_j][1]+2*dp[bjb_j][2]です。
...とはなりません。この場合、例えば次のような木を考えてみましょう。
この木において辺ibji-b_jを切断してしまうと、頂点iが、文字aとb両方含むという条件を満たさなくなります。よって、iを含む連結成分の文字がaだけになる通り数を引く必要があります。これはdp[i][0]そのものであるので、この値を後から引いてあげます。
 これらをすべての子頂点に対して行い、遷移に必要な情報が集まったとします。必要なのは辺の削除の通り数なので、すべての子頂点の情報の積を取ってやればよいことが分かります。よって、公式解説にもあるように、頂点iとその子頂点bjb_jに対して、dp[i][0],dp[i][1],dp[i][2]は次の式で更新できます。
dp[i][0]=(dp[bj][0]+dp[bj][2])dp[i][0] = \prod (dp[b_j][0] + dp[b_j][2])
dp[i][1]は更新不要dp[i][1]は更新不要
dp[i][2]=(dp[bj][0]+dp[bj][1]+2dp[bj][2])dp[i][0]dp[i][2] = \prod (dp[b_j][0] + dp[b_j][1] + 2 * dp[b_j][2]) - dp[i][0]

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

 頂点iにaが書かれている場合と、aとbを入れ替えて読み替えたものになります。

実装

 葉であるときのbase case の場合分けは面倒なので、m1 := (dp[i][0]またはdp[i][1]に代入して更新するための値)、m2 := (dp[i][2]に代入して更新するための値)という変数を作り、初期値を1にしてしまいましょう。
'''cpp:sample.cpp
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;
'''

提出コード

終わりに

 この記事が誰かの理解の助けになればうれしいです。説明の矛盾や誤字脱字などあったら教えてくれると助かります。ここまで読んでいただきありがとうございました。

Discussion

コメントにはログインが必要です。