K-Set Tree
Contest: CodeCraft-22 and Codeforces Round 795 (Div. 2)
Difficulty rating: 2500
Tags: combinatorics, dfs and similar, dp, math, trees

You are given a tree $$$G$$$ with $$$n$$$ vertices and an integer $$$k$$$. The vertices of the tree are numbered from $$$1$$$ to $$$n$$$.

For a vertex $$$r$$$ and a subset $$$S$$$ of vertices of $$$G$$$, such that $$$|S| = k$$$, we define $$$f(r, S)$$$ as the size of the smallest rooted subtree containing all vertices in $$$S$$$ when the tree is rooted at $$$r$$$. A set of vertices $$$T$$$ is called a rooted subtree, if all the vertices in $$$T$$$ are connected, and for each vertex in $$$T$$$, all its descendants belong to $$$T$$$.

You need to calculate the sum of $$$f(r, S)$$$ over all possible distinct combinations of vertices $$$r$$$ and subsets $$$S$$$, where $$$|S| = k$$$. Formally, compute the following: $$$$$$\sum_{r \in V} \sum_{S \subseteq V, |S| = k} f(r, S),$$$$$$ where $$$V$$$ is the set of vertices in $$$G$$$.

Output the answer modulo $$$10^9 + 7$$$.

Input format:
The first line contains two integers $$$n$$$ and $$$k$$$ ($$$3 \le n \le 2 \cdot 10^5$$$, $$$1 \le k \le n$$$).

Each of the following $$$n - 1$$$ lines contains two integers $$$x$$$ and $$$y$$$ ($$$1 \le x, y \le n$$$), denoting an edge between vertex $$$x$$$ and $$$y$$$.

It is guaranteed that the given edges form a tree.

Output format:
Print the answer modulo $$$10^9 + 7$$$.

Constraints:
Time limit: 3.0 s
Memory limit: 512.0 MB

Note:
The tree in the second example is given below:

We have $$$21$$$ subsets of size $$$2$$$ in the given tree. Hence, $$$$$$S \in \left\{\{1, 2\}, \{1, 3\}, \{1, 4\}, \{1, 5\}, \{1, 6\}, \{1, 7\}, \{2, 3\}, \{2, 4\}, \{2, 5\}, \{2, 6\}, \{2, 7\}, \{3, 4\}, \{3, 5\}, \{3, 6\}, \{3, 7\}, \{4, 5\}, \{4, 6\}, \{4, 7\}, \{5, 6\}, \{5, 7\}, \{6, 7\} \right\}.$$$$$$ And since we have $$$7$$$ vertices, $$$1 \le r \le 7$$$. We need to find the sum of $$$f(r, S)$$$ over all possible pairs of $$$r$$$ and $$$S$$$.

Below we have listed the value of $$$f(r, S)$$$ for some combinations of $$$r$$$ and $$$S$$$.

- $$$r = 1$$$, $$$S = \{3, 7\}$$$. The value of $$$f(r, S)$$$ is $$$5$$$ and the corresponding subtree is $$$\{2, 3, 4, 6, 7\}$$$.
- $$$r = 1$$$, $$$S = \{5, 4\}$$$. The value of $$$f(r, S)$$$ is $$$7$$$ and the corresponding subtree is $$$\{1, 2, 3, 4, 5, 6, 7\}$$$.
- $$$r = 1$$$, $$$S = \{4, 6\}$$$. The value of $$$f(r, S)$$$ is $$$3$$$ and the corresponding subtree is $$$\{4, 6, 7\}$$$.

Example input:
3 2
1 2
1 3

Example output:
25

Example input:
7 2
1 2
2 3
2 4
1 5
4 6
4 7

Example output:
849

Editorial:
Our task is to calculate $$$$$$\sum_{R \in V} \sum_{S \subseteq V, |S| = k} f(R, S)$$$$$$ over all possible $$$R, S$$$. Let us first focus on all possible $$$S$$$ just for a particular root $$$R = r$$$ i.e. let us root our tree at $$$r$$$ and try to find the sum $$$\sum_{S \subseteq V, |S| = k} f(R,S)$$$ just for that root.
We will calculate our answer using dynamic programming over the trees. In this technique, we will calculate some properties for each sub-tree and eventually get those properties for the entire tree.
The first property that we want to calculate for each sub-tree with node $$$v$$$ as the sub-tree root is — $$$cnt(v)$$$ which is the number of subsets of size $$$k$$$ such that sub-tree of $$$v$$$ is the minimum-size sub-tree covering it entirely. This can be calculated using combinatorics — first we calculate the total number of subsets of size $$$k$$$ in this sub-tree and then from it we can subtract the number of subsets of size $$$k$$$ which don't have sub-tree of $$$v$$$ as the minimum size sub-tree. $$$$$$cnt(v) = {size(v) \choose k} - \sum \limits_{u \in \text{children}} {size(u) \choose k}$$$$$$ Here, $$$u$$$ are all the children of $$$v$$$ and $$$size(x)$$$ represents the size of the sub-tree of $$$x$$$.
This first property that we calculated is very important for us:
• If we take the sum of $$$cnt(v)$$$ over every node, we will get the total number of subsets of size $$$k$$$.
• When the tree is rooted at $$$r$$$, $$$cnt(v)$$$ represents the number of subsets where sub-tree of $$$v$$$ is the smallest sub-tree containing a set $$$S$$$ with $$$k$$$ vertices.
• Conclusively, $$$f(R = r, S) = \sum \limits_{v = 1}^{n} cnt(v) \times size(v)$$$.
The second property that we want to find for each sub-tree is the $$$size(v)$$$ — the size of the sub-tree of $$$v$$$.
The third property that we want to find for each sub-tree is $$$cntsz(v) = cnt(v) \times size(v)$$$.
Now, we have $$$f(R = r, S)$$$ (as explained above) i.e. the contribution to the final answer when the root of the entire tree is fixed at $$$r$$$. We can calculate the final answer by fixing other nodes as roots and then summing these value up.
Notice what happens when we try to change the root from $$$r$$$ to one of it's children. The properties that we calculated for each sub-tree remain the same except for the old root and the new root. We can recalculate the properties for these two nodes using some clever arithmetic and get the new answer with a new root. This is known as re-rooting technique. The method to calculate the new properties are:
Note: We use $$$OR$$$ to represent Old Root. and $$$NR$$$ to represent New Root.
• $$$size_{new}(OR) = size_{old}(OR) - size_{old}(NR)$$$. (Subtracting the size of this branch.)
• $$$size_{new}(NR) = n$$$. (This is the main root.)
• $$$cnt_{new}(OR) = cnt_{old}(OR) - {size_{old}(OR) \choose k} + {size_{new}(OR) \choose k} + {size_{old}(NR) \choose k}$$$. (Removing contribution of old size and putting contribution of new size. Removing contribution of the branch.)
• $$$cnt_{new}(NR) = cnt_{old}(NR) - {size_{old}(NR) \choose k} + {size_{new}(NR) \choose k} - {size_{new}(OR) \choose k}$$$. (Removing contribution of old size and putting contribution of new size. Putting contribution of new brach.)
• $$$cntsz_{new}(OR) = cnt_{new}(OR) \times size_{new}(OR)$$$ (By definition.)
• $$$cntsz_{new}(NR) = cnt_{new}(NR) \times size_{new}(NR)$$$ (By definition.)
• $$$ans_{new} = ans_{old} - cntsz_{old}(OR) - cntsz_{old}(NR) + cntsz_{new}(OR) + cntsz_{new}(NR)$$$ (Subtracting old contribution and adding new contribution.)
The final answer is given by: $$$$$$finalans = \sum \limits_{v=1}^n ans_v$$$$$$