Codeforces darkshadows's blog 「DP on Trees Tutorial」Problem 3

Codeforces darkshadows's blog 「DP on Trees Tutorial」Problem 2の続編です。
codeforces.com
問題3.
「頂点数nの木が与えられたとき、頂点がk個以下のサブツリーの数はいくつあるか?」
ここでいう「サブツリー(sub tree)」は、今まで使っていた「部分木(subtree)」とは異なり、「k個の頂点が連結していればよい」ということに注意してください。
f:id:unosss:20200330233114p:plain

木全体の根を頂点1としておき、頂点Vを根とする部分木をS(V)としておきます。

まずは、サブツリーの総数を計算しましょう。
頂点Vを含む部分木S(V)のサブツリーの数をf(V)とします。
f(V)の計算方法を考えます。頂点Vの子の頂点をuとすると、

①頂点uを根とする部分木S(u)について、頂点uを含める場合はもちろんf(u)個のサブツリーがあります。
②頂点uを含めない場合でも頂点Vを含むサブツリーを作ることができます。この場合、「頂点uを含めない」という意味で1通りですね。
f:id:unosss:20200330234542p:plain

よって、f(V)は

f(V)=Π(i=1~n){ f(vi)+1 }

と計算できます。(Πは∑の掛け算バージョンです)
しかし、これではまだ不十分ですね。f(1)で数えているのは「頂点1を含むサブツリーの数」です。よって、頂点Vを含まないサブツリーの数であるg(V)を定義する必要があります。g(V)を計算するとき、子viを含むサブツリーの数f(vi)と子viを含まないサブツリーの数g(vi)を足せば良いですね。

g(V)=∑(i=1~n){ f(vi)+g(vi) }

と計算できます。これらを再帰的に計算していけば、f(1)+g(1)が答えとなります。

さて、元の問題に戻りましょう。頂点がk個以下のサブツリーの総数でしたね。
「頂点がk個以下」という条件が加わったので、

f(V,k):=頂点がk個以下で、頂点Vを含むサブツリーの総数

と定義しましょう。頂点Vの子をv1, v2, ... , vnとします。頂点Vを根とする、頂点がk+1個のサブツリーを形成することを考えます。
頂点viの部分木S(vi)のうち、ai個の頂点を使っているとすると、

k=∑(i=1~n){ ai }

となりますね。よって、

f(V, k+1)=Π(i=1~n){ f(vi, ai) }

と計算できます(もちろん、ありうるすべての組(a1, a2, ... , an)について計算する必要があります)。
これを計算するために、頂点Vについてdp1を次のように定義します。

dp1(i, j):=(頂点Vの子をv1, v2, ... , vnとしたとき、これらの部分木S(v1), S(v2), ... , S(vi)からj個の頂点を選ぶ方法)

すると、

dp1(i, j)=∑(m=0~k){ dp1(i-1, j-k)×f(i, k) }

と計算できます。頂点Vについて、

f(V, k)=dp1(n, k)

が計算でき、最終的な答えは、

∑(i=1~k){ f(V, i) }

をすべての頂点について計算すれば求まります。

f[N][K+1]

void rec(int cur_node){

     f[cur_node][1]=1
     dp_buffer[K] = {0} //f(V, 0~k)の計算結果を保管する。
     dp_buffer[0] = 1

     for(all v such that v is children of cur_node)
      rec(v)

      dp_buffer1[K] = {0} //dp1(i, j)を計算するために、毎回初期化されたものを使う。          
      for i=0 to K:
           for j=0 to K-i:
                dp_buffer1[i + j] += dp_buffer[i]*f[v][j]

      dp_buffer = dp_buffer1

     f[cur_node] = dp_buffer
}