Koxia and Tree (Codeforces 1770E)
Tags: Math, probabilities

Imi has an undirected tree with $$$n$$$ vertices where edges are numbered from $$$1$$$ to $$$n-1$$$. The $$$i$$$-th edge connects vertices $$$u_i$$$ and $$$v_i$$$. There are also $$$k$$$ butterflies on the tree. Initially, the $$$i$$$-th butterfly is on vertex $$$a_i$$$. All values of $$$a$$$ are pairwise distinct.Koxia plays a game as follows:  For $$$i = 1, 2, \dots, n - 1$$$, Koxia set the direction of the $$$i$$$-th edge as $$$u_i \rightarrow v_i$$$ or $$$v_i \rightarrow u_i$$$ with equal probability.  For $$$i = 1, 2, \dots, n - 1$$$, if a butterfly is on the initial vertex of $$$i$$$-th edge and there is no butterfly on the terminal vertex, then this butterfly flies to the terminal vertex. Note that operations are sequentially in order of $$$1, 2, \dots, n - 1$$$ instead of simultaneously.  Koxia chooses two butterflies from the $$$k$$$ butterflies with equal probability from all possible $$$\frac{k(k-1)}{2}$$$ ways to select two butterflies, then she takes the distance$$$^\dagger$$$ between the two chosen vertices as her score. Now, Koxia wants you to find the expected value of her score, modulo $$$998\,244\,353^\ddagger$$$.$$$^\dagger$$$ The distance between two vertices on a tree is the number of edges on the (unique) simple path between them.$$$^\ddagger$$$ Formally, let $$$M = 998\,244\,353$$$. It can be shown that the answer can be expressed as an irreducible fraction $$$\frac{p}{q}$$$, where $$$p$$$ and $$$q$$$ are integers and $$$q \not \equiv 0 \pmod{M}$$$. Output the integer equal to $$$p \cdot q^{-1} \bmod M$$$. In other words, output such an integer $$$x$$$ that $$$0 \le x < M$$$ and $$$x \cdot q \equiv p \pmod{M}$$$.

Input: The first line contains two integers $$$n$$$, $$$k$$$ ($$$2 \leq k \leq n \leq 3 \cdot {10}^5$$$) — the size of the tree and the number of butterflies.The second line contains $$$k$$$ integers $$$a_1, a_2, \dots, a_k$$$ ($$$1 \leq a_i \leq n$$$) — the initial position of butterflies. It's guaranteed that all positions are distinct.The $$$i$$$-th line in following $$$n − 1$$$ lines contains two integers $$$u_i$$$, $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$, $$$u_i \neq v_i$$$) — the vertices the $$$i$$$-th edge connects.It is guaranteed that the given edges form a tree.

Output: Output a single integer — the expected value of Koxia's score, modulo $$$998\,244\,353$$$.

Note: In the first test case, the tree is shown below. Vertices containing butterflies are noted as bold.  There are only $$$2$$$ butterflies so the choice of butterflies is fixed. Let's consider the following $$$4$$$ cases:  Edges are $$$1 \rightarrow 2$$$ and $$$2 \rightarrow 3$$$: butterfly on vertex $$$1$$$ moves to vertex $$$2$$$, but butterfly on vertex $$$3$$$ doesn't move. The distance between vertices $$$2$$$ and $$$3$$$ is $$$1$$$.  Edges are $$$1 \rightarrow 2$$$ and $$$3 \rightarrow 2$$$: butterfly on vertex $$$1$$$ moves to vertex $$$2$$$, but butterfly on vertex $$$3$$$ can't move to vertex $$$2$$$ because it's occupied. The distance between vertices $$$2$$$ and $$$3$$$ is $$$1$$$.  Edges are $$$2 \rightarrow 1$$$ and $$$2 \rightarrow 3$$$: butterflies on both vertex $$$1$$$ and vertex $$$3$$$ don't move. The distance between vertices $$$1$$$ and $$$3$$$ is $$$2$$$.  Edges are $$$2 \rightarrow 1$$$ and $$$3 \rightarrow 2$$$: butterfly on vertex $$$1$$$ doesn't move, but butterfly on vertex $$$3$$$ move to vertex $$$2$$$. The distance between vertices $$$1$$$ and $$$2$$$ is $$$1$$$. Therefore, the expected value of Koxia's score is $$$\frac {1+1+2+1} {4} = \frac {5} {4}$$$, which is $$$748\,683\,266$$$ after modulo $$$998\,244\,353$$$.In the second test case, the tree is shown below. Vertices containing butterflies are noted as bold. The expected value of Koxia's score is $$$\frac {11} {6}$$$, which is $$$831\,870\,296$$$ after modulo $$$998\,244\,353$$$.  

Sample Inputs:
3 2
1 3
1 2
2 3

5 3
3 4 5
1 2
1 3
2 4
2 5

Sample Outputs:
748683266

831870296

Solution (Python):
import io, os
from collections import deque
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
M = 998244353

def fastfrac(a, b):
    numb = pow(b, M - 2, M)
    return a % M * (numb % M) % M

def main(t):
    (n, k) = map(int, input().split())
    points = list(map(int, input().split()))
    for i in range(k):
        points[i] -= 1
    points = set(points)
    neigh = [[] for _ in range(n)]
    edges = []
    for _ in range(n - 1):
        (u, v) = map(int, input().split())
        neigh[u - 1].append(v - 1)
        neigh[v - 1].append(u - 1)
        edges.append((u - 1, v - 1))
    subnum = [0] * n
    subtot = [0] * n
    subroot = [0] * n
    parent = [-2] * n
    queue = deque()
    queue.append(0)
    parent[0] = -1
    children = [[] for _ in range(n)]
    seq = []
    while queue:
        index = queue.popleft()
        seq.append(index)
        for nextindex in neigh[index]:
            if parent[nextindex] > -2:
                continue
            parent[nextindex] = index
            children[index].append(nextindex)
            queue.append(nextindex)
    for index in seq[::-1]:
        if index in points:
            subnum[index] += 1
        for nextindex in neigh[index]:
            subtot[index] += subtot[nextindex] + subnum[nextindex]
            subnum[index] += subnum[nextindex]
    subroot[0] = subtot[0]
    for index in seq[1:]:
        subroot[index] = subroot[parent[index]] - subnum[index] + (k - subnum[index])
    prob = [0] * n
    invkk = fastfrac(1, k * k - k)
    inv2 = fastfrac(1, 2)
    ans = 0
    for index in range(n):
        if index in points:
            prob[index] = 1
            ans += subroot[index]
            ans %= M
    ans = ans * invkk % M
    for (u, v) in edges:
        (pu, pv) = (prob[u], prob[v])
        uv = pu * (1 - pv) * inv2 % M
        vu = pv * (1 - pu) * inv2 % M
        prob[u] = (pu * pv + uv + vu) % M
        prob[v] = prob[u]
        if parent[u] == v:
            (u, v) = (v, u)
            (uv, vu) = (vu, uv)
        ans += uv * (k - subnum[v] - 1 - subnum[v]) * 2 * invkk % M
        ans += vu * (subnum[v] - 1 - k + subnum[v]) * 2 * invkk % M
        ans %= M
    print(ans)
T = 1
t = 1
while t <= T:
    main(t)
    t += 1

Solution (Python):
import sys
from sys import stdin
from collections import deque

def NC_Dij(lis, start):
    ret = [float('inf')] * len(lis)
    ret[start] = 0
    q = deque([start])
    plis = [i for i in range(len(lis))]
    visit = []
    while len(q) > 0:
        now = q.popleft()
        visit.append(now)
        for nex in lis[now]:
            if ret[nex] > ret[now] + 1:
                ret[nex] = ret[now] + 1
                plis[nex] = now
                q.append(nex)
    return (ret, plis, visit)

def inverse(a, mod):
    return pow(a, mod - 2, mod)
mod = 998244353
half = inverse(2, mod)
uv = []
(n, k) = map(int, stdin.readline().split())
a = list(map(int, stdin.readline().split()))
lis = [[] for i in range(n)]
for i in range(n - 1):
    (u, v) = map(int, stdin.readline().split())
    u -= 1
    v -= 1
    lis[u].append(v)
    lis[v].append(u)
    uv.append((u, v))
(dlis, plis, visit) = NC_Dij(lis, 0)
a_under = [0] * n
dp = [0] * n
for i in range(k):
    a[i] -= 1
    a_under[a[i]] += 1
    dp[a[i]] += 1
visit.reverse()
for v in visit:
    if v != 0:
        a_under[plis[v]] += a_under[v]
ans_original = 0
for i in range(1, n):
    ans_original += (k - a_under[i]) * a_under[i]
    ans_original %= mod
ans = ans_original
for (u, v) in uv:
    if plis[u] != v:
        (u, v) = (v, u)
    assert plis[u] == v
    newv = 0
    newu = 0
    u_to_v = (1 - dp[v]) * dp[u] * half % mod
    v_to_u = (1 - dp[u]) * dp[v] * half % mod
    sizeU = a_under[u]
    sizeV = k - a_under[u]
    ans += (sizeU - 1 - sizeV) * u_to_v
    ans += (sizeV - 1 - sizeU) * v_to_u
    ans %= mod
    newv += ((1 - dp[v]) * dp[u] + dp[v]) * half
    newu += dp[u] * dp[v] * half
    newv += dp[u] * dp[v] * half
    newu += ((1 - dp[u]) * dp[v] + dp[u]) * half
    dp[u] = newu % mod
    dp[v] = newv % mod
ans *= inverse(k * (k - 1) // 2, mod)
print(ans % mod)

Solution (Python):
import sys
input = sys.stdin.readline
MOD = 998244353
INV2 = MOD // 2 + 1
(n, k) = map(int, input().strip().split())
p = [0] * (n + 1)
s = [0] * (n + 1)
for u in map(int, input().strip().split()):
    s[u] = 1
    p[u] = 1
e = []
g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    (u, v) = map(int, input().strip().split())
    e.append((u, v))
    g[u].append(v)
    g[v].append(u)
f = [0] * (n + 1)
q = [1]
ql = 0
while ql < len(q):
    u = q[ql]
    ql += 1
    for v in g[u]:
        if v != f[u]:
            q.append(v)
            f[v] = u
for u in q[n:1:-1]:
    s[f[u]] += s[u]
ans = 0
for (u, v) in e:
    if f[u] == v:
        (u, v) = (v, u)
    ans += (((p[u] - p[v]) * (k - 2 * s[v]) - p[u] - p[v] + 2 * p[u] * p[v]) % MOD * INV2 + s[v] * (k - s[v])) % MOD
    p[u] = (p[u] + p[v]) * INV2 % MOD
    p[v] = p[u]
print(ans * pow(k * (k - 1) // 2, MOD - 2, MOD) % MOD)