Diameter Cuts
Contest: Educational Codeforces Round 106 (Rated for Div. 2)
Difficulty rating: 2400
Tags: combinatorics, dfs and similar, dp, trees

You are given an integer $$$k$$$ and an undirected tree, consisting of $$$n$$$ vertices.

The length of a simple path (a path in which each vertex appears at most once) between some pair of vertices is the number of edges in this path. A diameter of a tree is the maximum length of a simple path between all pairs of vertices of this tree.

You are about to remove a set of edges from the tree. The tree splits into multiple smaller trees when the edges are removed. The set of edges is valid if all the resulting trees have diameter less than or equal to $$$k$$$.

Two sets of edges are different if there is an edge such that it appears in only one of the sets.

Count the number of valid sets of edges modulo $$$998\,244\,353$$$.

Input format:
The first line contains two integers $$$n$$$ and $$$k$$$ ($$$2 \le n \le 5000$$$, $$$0 \le k \le n - 1$$$) — the number of vertices of the tree and the maximum allowed diameter, respectively.

Each of the next $$$n-1$$$ lines contains a description of an edge: two integers $$$v$$$ and $$$u$$$ ($$$1 \le v, u \le n$$$, $$$v \neq u$$$).

The given edges form a tree.

Output format:
Print a single integer — the number of valid sets of edges modulo $$$998\,244\,353$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 256.0 MB

Note:
In the first example the diameter of the given tree is already less than or equal to $$$k$$$. Thus, you can choose any set of edges to remove and the resulting trees will have diameter less than or equal to $$$k$$$. There are $$$2^3$$$ sets, including the empty one.

In the second example you have to remove the only edge. Otherwise, the diameter will be $$$1$$$, which is greater than $$$0$$$.

Here are the trees for the third and the fourth examples:

Example input:
4 3
1 2
1 3
1 4

Example output:
8

Example input:
2 0
1 2

Example output:
1

Editorial:
The task is obviously solved by dynamic programming, so our first reaction should be to start looking for meaningful states for it. Obviously, one of the states is the vertex which subtree we are processing. We can choose the root for the tree arbitrarily, let it be vertex $$$1$$$. What can be the other helpful state?
Consider the method to find the diameter of the subtree of vertex $$$v$$$. The diameter can be one of the following paths: either the longest path that is completely in some subtree of $$$v$$$ or the concatenation of the longest paths that start in vertex $$$v$$$ and end in different subtrees.
The diameter is the longest path. Thus, the diameter being less than or equal to $$$k$$$ means that all paths should have length less than or equal to $$$k$$$.
If we can guarantee that no path that is completely in some subtree of $$$v$$$ have length greater than $$$k$$$, then we will only have to worry about not concatenating long paths from different subtrees. Phrase it the other way around: if we never concatenate the paths from the different subtrees in such a way that their total length is greater than $$$k$$$, then no diameter will be greater than $$$k$$$.
Thus, we can attempt to have $$$dp[v][len]$$$ — the number of ways to cut some edges in the subtree of $$$v$$$ in such a way that there is no path of length greater than $$$k$$$ and the longest path starting at vertex $$$v$$$ has length $$$len$$$.
Now for the transitions. For the simplicity, let vertex $$$v$$$ have exactly two children. It's not too hard to merge their $$$dp$$$'s. Iterate over the length $$$i$$$ of the first child, the length $$$j$$$ of the second child. If $$$(i+1)+(j+1) \le k$$$, then you can concatenate their longest paths and the longest path for $$$v$$$ will be of length $$$max(i,j)+1$$$. You can also cut either of the edges from $$$v$$$ to the first child or to the second child.
The approach is good, however, it's not clear how to make it work on a larger number of children. Also, the complexity sounds pretty bad.
Instead of merging children to each other, let's merge each child to the $$$dp$$$ of $$$v$$$ one by one. $$$dp[v][len]$$$ can store the current maximum length over all processed children. When processing a new child, you can choose to cut or not to cut the edge to it. So you can iterate over the current longest path from $$$v$$$ and the longest path from that child.
So far, the only way to estimate the complexity is to say that each child has to merge its dp to the parent in $$$O(k^2)$$$, thus making the algorithm $$$O(nk^2)$$$. That's obviously too slow.
The trick that makes the solution fast is to iterate not to $$$k$$$ but to the height of the subtree of $$$v$$$ and the subtree of a child. Surely, that is allowed, since the path just can't grow longer than that value.
Consider the even worse option: not the height but the size of the subtree. It's easy to see that the size is always greater or equal than the height. Interpret the merge the following way: enumerate the vertices inside all the subtrees of the processed children and the vertices inside the subtree of the new child. Iterating up to the size of the subtree is the same number of moves as going over the vertices in it. The merge will go over all the pairs of vertices such that the first vertex of the pair is in the first set and the second vertex is in the second set. Thus, each pair of vertices of the tree will be processed exactly once (in lca of these vertices). There are $$$O(n^2)$$$ such pairs, thus, such $$$dp$$$'s work in $$$O(n^2)$$$.
Overall complexity: $$$O(n^2)$$$.