K Paths
Contest: Avito Code Challenge 2018
Difficulty rating: 3100
Tags: combinatorics, data structures, dp, fft, math

You are given a tree of $$$n$$$ vertices. You are to select $$$k$$$ (not necessarily distinct) simple paths in such a way that it is possible to split all edges of the tree into three sets: edges not contained in any path, edges that are a part of exactly one of these paths, and edges that are parts of all selected paths, and the latter set should be non-empty.

Compute the number of ways to select $$$k$$$ paths modulo $$$998244353$$$.

The paths are enumerated, in other words, two ways are considered distinct if there are such $$$i$$$ ($$$1 \leq i \leq k$$$) and an edge that the $$$i$$$-th path contains the edge in one way and does not contain it in the other.

Input format:
The first line contains two integers $$$n$$$ and $$$k$$$ ($$$1 \leq n, k \leq 10^{5}$$$) — the number of vertices in the tree and the desired number of paths.

The next $$$n - 1$$$ lines describe edges of the tree. Each line contains two integers $$$a$$$ and $$$b$$$ ($$$1 \le a, b \le n$$$, $$$a \ne b$$$) — the endpoints of an edge. It is guaranteed that the given edges form a tree.

Output format:
Print the number of ways to select $$$k$$$ enumerated not necessarily distinct simple paths in such a way that for each edge either it is not contained in any path, or it is contained in exactly one path, or it is contained in all $$$k$$$ paths, and the intersection of all paths is non-empty.

As the answer can be large, print it modulo $$$998244353$$$.

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

Note:
In the first example the following ways are valid：

- $$$((1,2), (1,2))$$$,
- $$$((1,2), (1,3))$$$,
- $$$((1,3), (1,2))$$$,
- $$$((1,3), (1,3))$$$,
- $$$((1,3), (2,3))$$$,
- $$$((2,3), (1,3))$$$,
- $$$((2,3), (2,3))$$$.

In the second example $$$k=1$$$, so all $$$n \cdot (n - 1) / 2 = 5 \cdot 4 / 2 = 10$$$ paths are valid.

In the third example, the answer is $$$\geq 998244353$$$, so it was taken modulo $$$998244353$$$, don't forget it!

Example input:
3 2
1 2
2 3

Example output:
7

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

Example output:
10

Editorial:
For each vertex $$$v$$$ let's calculate $$$f_v$$$ — number of ways to choose $$$k$$$ ends for paths in their subtree, such that these paths are edge-intersecting, and $$$g_v$$$ — sum of $$$f_u$$$ by all vertices $$$u$$$ in subtree of $$$v$$$. Let $$$sz_v$$$ will be number of vertices in subtree of $$$v$$$.
Then let $$$P(x) = (1 + x \cdot sz_{u_{0}}) \cdot (1 + x \cdot sz_{u_{1}}) \ldots \cdot (1 + x \cdot sz_{u_{deg - 1}})$$$, (where $$$u_0, u_1, \ldots, u_{deg - 1}$$$ are the childs of vertex $$$v$$$), $$$ = \sum(a_k \cdot x^k)$$$
How to calculate $$$P(x)$$$? We will use divide and conquer, let's recursively calculate answer for first half of element, for second, and then multiplicate.
$$$T(n) = 2 T(n / 2) + O(n log n) = O(n log^2 n)$$$
Note, that degree of polynom is not $$$n$$$, but $$$deg_v$$$!.
Then $$$f_v = \sum(a_x \cdot C_k^x \cdot x!)$$$.
Then for fixed non-vertical path $$$u \to v$$$, number of ways to choose $$$k$$$ paths, such that their intersection is $$$u \to v$$$ is $$$f_u \cdot f_v$$$.
We can sum this up for each non-vertical path with simple tree dp.
So we are working in $$$O(deg_v log^2 n)$$$ for fixed $$$v$$$, summing up for each vertex we can get that this part of solution is working in $$$O(n log^2 n)$$$.
Let fix one end of vertical path $$$v$$$ and it's son $$$u$$$, such that second end of vertical path is inside his subtree, Then let $$$Q(x) = P(x) \cdot(1 + x \cdot (n - sz_v)) / (1 + x \cdot sz_u) = \sum(b_k \cdot x^k)$$$, when we need add $$$\sum(b_x \cdot C_k^x \cdot x!) \cdot g_u$$$ to answer.
Then we can note, that (for fixed $$$v$$$) number of distinct $$$sz_u$$$'s — $$$O(\sqrt{n})$$$
And for each $$$sz_u$$$ we can get $$$Q(x)$$$ in $$$O(deg_v)$$$, at first we need to calculate (one time!) $$$P(x) \cdot(1 + x \cdot (n - sz_v)) = \sum(c_k \cdot x^k)$$$, and then $$$b_0 = 1$$$, и $$$b_x = c_x - b_{x-1} \cdot sz_u$$$.
So, this part is working in $$$O(deg_v \cdot \sqrt{n})$$$ for fixed $$$v$$$, summing up for each vertices we can get that this part of solution is working in $$$O(n \sqrt{n})$$$, so we solved the task in $$$O(n log^2 n + n \sqrt{n})$$$
Bonus: solve in $$$n log^2 n$$$.
300iq's nsqrt isaf27's nlog^2