P8334 [ZJOI2022] Depth-First Search

Kanjou Karen is a girl who likes algorithms. Among many algorithms, she especially likes depth-first search (DFS). One day, Karen obtained a rooted tree with root $\mathit{root}$. Each node $x$ on the tree has a weight $a_x$. On a tree, starting from $x$ to search for node $y$, if we use depth-first search, it can be described by the following procedure: 1. Set the recursion stack to be empty. 2. First, push node $x$ into the recursion stack. 3. Pop the top node from the recursion stack. If this node is $y$, then stop the procedure. Otherwise, if there exists an unvisited direct child, then uniformly at random choose one child and push it into the recursion stack. 4. Repeat step 3 until there is no unvisited direct child. 5. Push the parent node into the recursion stack and repeat step 3. 6. Repeat step 5 until the current node is $x$, and the procedure ends. We define $f(x, y)$ to be valid if and only if $y$ is in the subtree of $x$. Its value is the expected minimum weight among all nodes visited (including $x$ and $y$) during the DFS on the subtree of $x$ starting from $x$ while searching for $y$. Karen wants to know, for all valid pairs $(x, y)$, the value of $\sum f(x, y)$. You only need to output the result modulo $998244353$. Specifically, if the answer in lowest terms is $\frac{a}{b}$, output $a \times b^{-1} \bmod 998244353$.

The input contains 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, \mathit{root}$, representing the size of the tree and the index of the root. The next line contains $n$ integers $a_1, a_2, \ldots, a_n$, representing the weight of each node. The next $n - 1$ lines each contain two integers $u, v$, indicating that there is an edge between $u$ and $v$.

For each test case, output one line containing one integer, representing $\sum f(x, y)$ over all valid pairs $(x, y)$ modulo $998244353$.

For all test points, it holds that $1 \le T \le 100$, $\sum n \le 8 \times {10}^5$, $1 \le n \le 4 \times {10}^5$, $1 \le \mathit{root}, u, v \le n$, $1 \le a_i \le {10}^9$. The detailed limits for each test point are shown in the table below: | Test Point ID | $\sum n \le$ | $n \le$ | Special Constraints | |:-:|:-:|:-:|:-:| | $1$ | $50$ | $10$ | None | | $2 \sim 4$ | $40000$ | $5000$ | None | | $5 \sim 10$ | $4 \times {10}^5$ | ${10}^5$ | None | | $11$ | $8 \times {10}^5$ | $4 \times {10}^5$ | The tree generation is random | | $12$ | $8 \times {10}^5$ | $4 \times {10}^5$ | The tree is a chain | | $13$ | $8 \times {10}^5$ | $4 \times {10}^5$ | The degree of the root is $n - 1$ | | $14 \sim 20$ | $8 \times {10}^5$ | $4 \times {10}^5$ | None | For test point $11$, the tree is generated as follows: take $1$ as the root. For each node $i \in [2, n]$, uniformly at random choose a node from $[1, i - 1]$ as its parent. Then the labels are randomly permuted. Translated by ChatGPT 5