鹿島建設プログラミングコンテスト2024（AtCoder Beginner Contest 340）G問題メモ

G - Leaf Color

$N$ 頂点の木があり、頂点 $i$ には色 $A_i$ が塗られている
以下の条件をともに満たす（空でない）誘導部分グラフの個数を $\mod{998244353}$ で求めよ
- 誘導部分グラフも木である
- 葉（次数が1の頂点）には全て同じ色が塗られている
$1 \le N \le 2 \times 10^5$

まず、頂点数1の木は条件を満たすが、ちょっと特殊なので分けて考える。とりあえず $N$ 個ある。

頂点数2以上の条件を満たす木は、以下で作れる。

    ④         (例) ④
   /｜\            /｜\
  ④○④     →   ④○④
 /| /\  \            \
④○○④ ○           ④

それ以外の頂点を繋げると、色④以外の葉ができてしまうので、条件を満たさない。
上記の方法で作れる木の個数を数えればよい。

「じゃあ、④が5個あったら、そこから2個以上選ぶ場合の数は $2^5-1-5=26$ 通りだ！」というと、そう単純でもない。

元の木で一直線上に並ぶ頂点は、端の頂点などを2つ選んだ場合、間の頂点は選んでも選ばなくても、できあがる木としては同じになってしまう。
「木としてユニークな」個数を数えなければならない。

ある同じ色の頂点の選び方において、その頂点を選ばないとできあがる木が変わってしまうような頂点を「有効頂点」とする。

    ④     ←④は選んでも選ばなくても、他が一緒ならできあがる木は一緒
   /｜\
  ④○❹   ←❹は有効頂点（できあがる木において葉になる頂点ともいえる）
 /   \
❹   ❹

木DPをする。

$DP[v,c]=v$ の部分木のみに着目した、色 $c$ の有効頂点の選び方の個数
- ただし、 $v$ ができあがる木に含まれ、そこから親にも繋がりうるような選び方のみを管理

更新を考える。
$DP[v]$ は $v$ の親に繋がる分のみを管理する。 $v$ 以下で完結するような選び方は、適宜答えを加算していくための変数 ans に足していき、DPには含めない。

    ：
    ①←v
  / ｜ \     DP[左の子] = {①: 2}
 ① ② ②    DP[中の子] = {①: 1, ②: 1}
 ｜ ｜       DP[右の子] = {②: 1}
 ① ①

(a) $v$ が次数1となり、子からつながるグラフを終了させる場合。

これは、 $v$ の色 $A_v$ について、各子の $DP[\cdot,A_v]$ を合計すればよい。
上記例では $v$ の色は①なので、左2 + 中1 = 3通り。
ansに加算し、 $DP[v]$ には含めない。

(b) これは単に $DP[v,A_v]+=1$

(d),(e) は、子を1つずつマージしていくことで求められる。いずれも同じ値になるので、答えに加算しつつ $DP[v]$ にも加算する。

まとめて、 $DP[v]$ および $v$ で完結する分のansへの加算は、以下のアルゴリズムで求められる。

これをそのままやるとTLEだが、マージテクにより、 $DP[v]$ と $DP[w]$ の大きい方に小さい方を追加していくと $O(N \log{N})$ で収まるようになる。