P10060 [SNOI2024] Tree V Diagram

You are given an unrooted tree with $n$ nodes, numbered $1, 2, \ldots, n$. Define the distance $\operatorname{dis}(i, j)$ between two nodes on the tree as the number of edges on the simple path from node $i$ to node $j$. Now there are $k$ key nodes $a_1, a_2, \ldots, a_k$. For each node, we want to find which key node is closest to it. That is, for a node $v$, let $f(v)$ be the index $i$ that minimizes $\operatorname{dis}(v, a_i)$. If multiple indices $i$ satisfy this, choose the smallest such $i$. Now you are given $f(1), f(2), \ldots, f(n)$. How many possible tuples $(a_1, a_2, \ldots, a_k)$ satisfy the condition? Since the answer may be large, output it modulo $998244353$.

Multiple test cases. The first line contains an integer $T$, the number of test cases. For each test case, the first line contains two integers $n, k$, the number of nodes and the number of key nodes. The next $n - 1$ lines each contain two integers $u, v$, representing an edge $(u, v)$ of the tree. The next line contains $n$ integers $f(1), f(2), \ldots, f(n)$. Note: the testdata is **not guaranteed** to have at least one valid tuple $(a_1, a_2, \ldots, a_k)$.

For each test case, output one integer: the answer modulo $998244353$.

**[Sample \#1 Explanation]** In the first sample, for the second test case, one solution is $(1, 2)$. For the third test case, there are two solutions: $(2, 1), (3, 1)$. Note that when multiple nodes have the same distance, we choose the smallest index $i$, not the smallest $a_i$. --- **[Sample \#3]** See the attached files `voronoi/voronoi3.in` and `voronoi/voronoi3.ans`. This sample satisfies the constraints of test points $3 \sim 4$. --- **[Sample \#4]** See the attached files `voronoi/voronoi3.in` and `voronoi/voronoi3.ans`. This sample satisfies the constraints of test points $7 \sim 10$. --- **[Constraints]** For all testdata, it is guaranteed that $1 \le T \le 10$, $2 \le k \le n \le 3 \times {10}^3$, and $1 \le f(i) \le k$. Details: | Test Point ID | $n \le$ | Special Property | |:-:|:-:|:-:| | $1 \sim 2$ | $15$ | None | | $3 \sim 4$ | $3000$ | A | | $5 \sim 6$ | $3000$ | B | | $7 \sim 10$ | $3000$ | None | Special Property A: the tree is guaranteed to be a chain. Special Property B: $k = 2$ is guaranteed. Translated by ChatGPT 5