This blog is about problem 3 from this blog.

Given a tree T of N nodes and an integer K, find number of different sub trees of size less than or equal to K.

This is a very useful problem in the whole world of cp. The solution given in that blog is O(n * K²) but we want to change it to O(n * K).

Consider Sub_v as the size of the subtree of vertex v.

The solution is similar to the given solution in that blog and the difference is very simple. You don't have to iterate over all K values. It's obvious that if i > Sub_v then dp[v][i] = 0. So each time you have to iterate over min(K, Sub_v) values.

The code becomes like this.

void solve(int v)
{
    Sub[v] = 1;
    dp[v][0] = 1;
    dp[v][1] = 1;

    for(auto u : childrens[v])
    {
       solve(u);

       fill(tmp , tmp + k + 1 , 0);
       for(int i = 1; i <= min(Sub[v] , k); i++)
         for(int j = 0; j <= Sub[u] && i + j <= K; j++)
          tmp[i + j] += dp[v][i] * dp[u][j];

                    Sub[v] += Sub[u];

       for(int i = 0; i <= min(K , Sub[v]); i++)
         dp[v][i] = tmp[i];
    }
}

The Proof Of The Complexity

First we introduce the array EnterTime. If EnterTime_v = x then v is the x'th vertex we have entered to solve. Every vertex has a st_v which is an array and in the very beginning st_v = {v} for all v. We'll change this set later and I'll explain how. Also consider st_v always sorted by EnterTime.

We want to build a digraph H of size n. When we're about to update dp[v] from its child u, add a directed edge from every vertex is st_u to every vertex in st_v in the digraph H.

Then update st_v by adding the first K - size(st_v) vertices from st_u. (If size(st_v) = k then ignore this) Also erase every vertex from st_u.

The complexity is equal to number of edges of H. We want to prove that number of edges in H is at most n * (2 * K) by proving that every vertex in H has at most 2 * K edges going out from it.

By our definition at any time, for any vertex v there's at most 1 vertex u such that v is in st_u. Consider some w's last set is st_u and we're about to update dp[v] from its child u. Consider w being in the p'th position in st_u. By our definition edges that go out of w are edges to the first p - 1 vertices in st_u.

We know that p ≤ K so until now, edges going out from w are at most K. When we're updating dp[v] we'll add at most K to that number. So by far w has at most 2 * K edges going out from it. Also we've considered st_u to be the last set that w is in. Hence there won't be any other edges going out from w in the future.

So we have proved number of edges in H is at most n * (2 * K) so the complexity is O(n * K).