Tree Cutting
Contest: Codeforces Round 854 by cybercats (Div. 1 + Div. 2)
Difficulty rating: 3200
Tags: bitmasks, dfs and similar, dp

You are given a tree with $$$n$$$ vertices.

A hero $$$k$$$ times do the following operation:

- Choose some edge.
- Remove it.
- Take one of the two remaining parts and delete it.
- Write the number of vertices in the remaining part.

You are given an initial tree and the a sequence of written numbers. Find the number of ways to make operations such that the written numbers are equal to the given numbers. Due to the answer can be big, find it by modulo $$$998\,244\,353$$$. Two ways are considered different, if on some operation edge or remaining part are selected differently.

Input format:
The first line contains a single integer $$$n$$$ ($$$2 \leq n \leq 5000$$$) — the number of vertices.

Each of the next $$$n-1$$$ lines contains two integers $$$s$$$, $$$f$$$ ($$$1 \leq s, f \leq n$$$, $$$s \neq f$$$) — description of edge $$$(s, f)$$$.

Next line contains a single integer $$$k$$$ ($$$1 \leq k \leq \min{(6, n - 1)}$$$) — the number of operations.

Next line contains $$$k$$$ integers $$$s_1, s_2, \ldots, s_k$$$ ($$$n > s_1 > s_2 > \ldots > s_k \geq 1$$$) — written numbers.

Output format:
Print a single integer — the answer to the problem by modulo $$$998\,244\,353$$$.

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

Note:
In the first test there are four possible ways to make operations:

- Remove the edge $$$(1, 2)$$$ and delete vertex $$$1$$$. Remove the edge $$$(2, 3)$$$ and delete vertex $$$2$$$.
- Remove the edge $$$(1, 2)$$$ and delete vertex $$$1$$$. Remove the edge $$$(3, 2)$$$ and delete vertex $$$3$$$.
- Remove the edge $$$(3, 2)$$$ and delete vertex $$$3$$$. Remove the edge $$$(1, 2)$$$ and delete vertex $$$1$$$.
- Remove the edge $$$(3, 2)$$$ and delete vertex $$$3$$$. Remove the edge $$$(2, 1)$$$ and delete vertex $$$2$$$.

In the second test there are two possible ways to make operations:

- Remove the edge $$$(4, 1)$$$ and delete the part with vertex $$$4$$$. Remove the edge $$$(2, 3)$$$ and delete the part with vertex $$$2$$$.
- Remove the edge $$$(4, 1)$$$ and delete the part with vertex $$$4$$$. Remove the edge $$$(3, 2)$$$ and delete the part with vertex $$$3$$$.

Example input:
3
1 2
2 3
2
2 1

Example output:
4

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

Example output:
2

Editorial:
We should calculate the number of ways to choose some $$$k$$$ edges of our tree (and directions of them) corresponding to operations, such that the operations that will be done with them will result in the given sequence of written numbers.
To calculate these ways let's consider a subtree $$$dp_{v,mask}$$$ for our main tree. Here $$$v$$$ corresponds to subtree of our tree with root in vertex $$$v$$$, also we include an edge from $$$v$$$ to the parent into this subtree. $$$mask$$$ is a submask of $$$\{1, 2, \ldots, k\}$$$, meaning that we have chosen in this subtree edges corresponding to operations from $$$mask$$$.
First of all, consider the case, when the edge from $$$v$$$ to parent is not chosen to any of the operations. So let's iterate over children of $$$v$$$ and calculate the same dynamic programming for the considered prefix of subtrees. To add the new subtree, we should iterate over $$$mask_1$$$ mask of operations that have been done on prefix and $$$mask_2$$$ mask of operations that have been done in new subtree. For these masks the condition $$$mask_1 \cap mask_2 = \emptyset$$$ should hold.
In the case, when the edge from $$$v$$$ to parent is chosen let's iterate of the operation $$$i$$$ corresponding to it. After that, $$$sz_{v} = s_i$$$ or $$$sz_{v} = n - s_i$$$ should hold. If $$$sz_{v} = s_i$$$ we can choose this edge for the operation $$$i$$$ and orient it from parent to $$$v$$$. So, the operations in $$$mask \setminus \{i\}$$$ should be with indices $$$> i$$$. Iterate such masks and update the $$$dp$$$ value corresponding to it (using dp values already calculated for subtree of $$$v$$$ for the case when edge to parent is not chosen). In the case $$$sz_{v} = n - s_i$$$ similarly $$$sz_{v} = n - s_i$$$ should hold and operations in $$$mask \setminus \{i\}$$$ should be with indices $$$\leq i$$$.
We can see, that the number of ways calculated with such dp is correct, because if we assigned operation $$$i$$$ to some edge with some direction, we ensured the written size after the operation is correct and that the operations in the subtree have been done before or after the operation $$$i$$$ (in relation to direction of edge).
Time complexity: $$$O(n (3^k + 2^k k))$$$, but it can be improved to $$$O(n 2^k k)$$$ with sum over subsets calculation.