競プロ典型90問 073 We Need Both a and b(★5)
木上の数え上げ問題であるので, DPをしていく.
$ dp_{i, j} := 頂点$ iの部分木であって, 状態が$ jであるものの場合の数 と定義する. ここで状態$ 0 := 'a'しかない, $ 1 := 'b'しかない, $ 2 := 'a'と'b'の両方ある とおく.
頂点$ iに'a'が書かれた場合と, 'b'が書かれた場合に場合分けして考える.
頂点$ iに'a'が書かれた場合
$ dp_{i, 0}を求める. 頂点$ iの子の部分木について, 状態$ 0であるものは辺を切断しないとすると含めてよく, 状態$ 2であるものは辺を切断するとすると含められるから, $ dp_{i, 0} = (dp_{nv_1, 0} + dp_{nv_1, 2}) \cdot (dp_{nv_2, 0} + dp_{nv_2, 2}) \cdot ... \cdot (dp_{nv_k, 0} + dp_{nv_k, 2}))となる.
$ dp_{i, 2}を求める. 頂点$ iの子の部分木について, 辺を削除する場合については状態$ 2のもののみ含めてよく, 辺を削除しない場合についてはすべての状態を含めてもよい. この場合の数から, $ dp_{i, 0}の分を除いて(包含関係にある)$ dp_{i, 2} = (dp_{nv_1, 0} + + dp_{nv_1, 1} + 2dp_{nv_1, 2}) \cdot (dp_{nv_2, 0} + dp_{nv_2, 1} + 2dp_{nv_2, 2}) \cdot ... \cdot (dp_{nv_k, 0} + dp_{nv_k, 1} + 2dp_{nv_k, 2}))となる.
頂点$ iに'b'が書かれた場合
上の場合と同様にして考えることにより遷移を導くことができる.
よって, 頂点$ 1からこの木DPをしていくことにより, $ O(N)でこの問題を解くことができた.