[1] Counting directed trees

First, let’s count the number of good trees.

Good trees correspond one-to-one with trees where the root has a degree of \(d_i\) and all other vertices have a degree of \(d_i + 1\). This counting is well-known (see ABC303-Ex) and the number can be written as follows:

\[\displaystyle\frac{(N-2)!}{\prod (e_1(i) - 1)!}.\]

(For convenience, we define the function \(e_r(i)\) as \(e_r(i)=d_i\) when \(i=r\), and \(e_r(i)=d_i+1\) when \(i\neq r\).)

[2] Tail wagging the dog

We will find the answer by summing the number of trees in which vertex \(v\) is a good vertex for all vertices. If \(v=1\) or \(d_v=0\), then \(v\) is always a good vertex, so we simply count the number of good trees. From now on, we consider the case where \(v\geq 2\) and \(d_v > 0\).

We fix the set of vertices in the subtree rooted at vertex \(v\), and call it \(S\). Then, from the relationship between the in-degree and out-degree, \(S\) must satisfy \(\displaystyle\sum_{v\in S} d_v = |S| - 1\). When this holds, the number of good trees is (the number of subtrees rooted at \(v\) composed of elements of \(S\)) \(\times\) (the number of trees composed of \(v\) and the elements not in \(S\)), which can be written as follows using [1]:

\[\displaystyle\frac{(|S|-2)!}{\prod_{i \in S} (e_v(i) - 1)!}\times \frac{(N-|S|+1 - 2)!}{\prod_{i \notin S} (e_1(i) - 1)!}.\]

This can be simplified as follows:

\[\displaystyle\frac{(|S|-2)!(N-|S| - 1)!d_1d_v}{ \prod d_i!}.\]

[3] Dynamic programming

We fixed \(S\) in [2], but since the derived formula depends only on \(v\) and \(|S|\), it can be calculated together using dynamic programming. Specifically, we can let \(\mathrm{dp}[i][j][k]\) be the number of ways to decide whether to include vertices \(N\) through \(i\) in \(S\) so that the sum of out-degrees is \(j\) and \(|S|=k\).

Since there are \(\mathrm{O}(N^3)\) states and \(\mathrm{O}(1)\) transitions for each state, we can solve this problem in \(\mathrm{O}(N^3)\).

(Above is a modification of a translation by GPT-4.)