次数 1 の頂点の色 $c$ を固定した時の問題を $O(N)$ で解く方法です。
次の動的計画法を考えます。
- $\mathrm{dp}[u]:=$ 頂点 $u$ を根とする部分木の中で、次数 1 の頂点の色が全て $c$ であるものの通り数。
- $\mathrm{dp}_2[u]:=$ 頂点 $u$ を根とする部分木の中で、次数 1 の頂点の色が全て $c$ であるものの通り数。ただし頂点 $u$ が色 $c$ でなく次数 1 であったとしても、頂点 $u$ が孤立点でなければよいものとする。
頂点 $u$ の色で場合分けをします。頂点 $u$ の子の集合を $\mathrm{child}(u)$ と表記します。
(1) 頂点 $u$ の色が $c$ である時
全ての $v \in \mathrm{child}(u)$ について、$v$ を根とする部分木を含めるか含めないかを考えると、
$\mathrm{dp}[u] = \prod_{v \in \mathrm{child}(u)} (\mathrm{dp}_2[v] + 1)$
$\mathrm{dp}_2[u] = \prod_{v \in \mathrm{child}(u)} (\mathrm{dp}_2[v] + 1)$
が成り立ちます。
(2) 頂点 $u$ の色が $c$ でない時
先程と同様に、全ての $v \in \mathrm{child}(u)$ について、$v$ を根とする部分木を含めるか含めないかを考えます。ただし、頂点 $u$ が次数 1 になってしまってはいけないので、頂点 $u$ の子ノードを根とする部分木を 2 つ以上含める必要があります。よって、
$\mathrm{dp}[u] = \prod_{v \in \mathrm{child}(u)} (\mathrm{dp}_2[v] + 1) - \sum_{v \in \mathrm{child}(u)} \mathrm{dp}_2[v] - 1$
$\mathrm{dp}_2[u] = \prod_{v \in \mathrm{child}(u)} (\mathrm{dp}_2[v] + 1) - 1$
が成り立ちます。（全体から子ノードを 1 つだけ採用するパターンと、子ノードを 1 つも採用しないパターンを引いています。）
また、$\mathrm{dp}_2$ に関しては、$u$ が孤立点でなければ、$u$ の親ノードが $u$ を部分木に含めることによって $u$ の次数が 2 以上になる可能性があります。
よって、頂点 $u$ の次数が 1 で確定してしまうのは現時点で頂点 $u$ が孤立点になってしまうパターンのみなので、子ノードを 1 つも採用しないパターンのみを除きます。
以上の動的計画法を行い、最終的に、$\sum_{u=1}^{N} \mathrm{dp}[u]$ が求める答えです。
