Sasha and a Walk in the City (Codeforces 1929D)
Tags: Math, trees

Sasha wants to take a walk with his girlfriend in the city. The city consists of $$$n$$$ intersections, numbered from $$$1$$$ to $$$n$$$. Some of them are connected by roads, and from any intersection, there is exactly one simple path$$$^{\dagger}$$$ to any other intersection. In other words, the intersections and the roads between them form a tree.Some of the intersections are considered dangerous. Since it is unsafe to walk alone in the city, Sasha does not want to visit three or more dangerous intersections during the walk.Sasha calls a set of intersections good if the following condition is satisfied: If in the city only the intersections contained in this set are dangerous, then any simple path in the city contains no more than two dangerous intersections.However, Sasha does not know which intersections are dangerous, so he is interested in the number of different good sets of intersections in the city. Since this number can be very large, output it modulo $$$998\,244\,353$$$.$$$^{\dagger}$$$A simple path is a path that passes through each intersection at most once.

Input: Each test consists of multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains a single integer $$$n$$$ ($$$2 \le n \leq 3 \cdot 10^5$$$) — the number of intersections in the city.The next $$$(n - 1)$$$ lines describe the roads. The $$$i$$$-th line contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$, $$$u_i \ne v_i$$$) — the numbers of the intersections connected by the $$$i$$$-th road.It is guaranteed that these roads form a tree.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$3 \cdot 10^5$$$.

Output: For each test case, output a single integer — the number of good sets of intersections modulo $$$998\,244\,353$$$.

Note: In the first test case, there are $$$2^3 = 8$$$ sets of intersections. All of them are good, except for the set $$$\{1, 2, 3\}$$$, because if intersections $$$1, 2$$$, and $$$3$$$ are dangerous, then the simple path $$$1 - 2 - 3$$$ contains $$$3$$$ dangerous intersections. Thus, there are $$$7$$$ good sets.In the second test case, there are $$$2^4 = 16$$$ sets of intersections. Among them, the sets $$$\{1, 2, 3, 4\}$$$, $$$\{1, 2, 3\}$$$, $$$\{1, 3, 4\}$$$, $$$\{2, 3, 4\}$$$ are not good. Thus, there are a total of $$$12$$$ good sets. The city layout is shown below:  

Sample Inputs:
4
3
1 3
3 2
4
3 4
2 3
3 1
5
1 2
3 4
5 1
2 3
4
1 2
2 3
3 4

Sample Outputs:
7
12
16
11

Solution (Python):
mod = 998244353
import sys
input = sys.stdin.readline
for _ in range(int(input())):
    n = int(input())
    g = [[] for k in range(n)]
    d = [0] * n
    for i in range(n - 1):
        (u, v) = map(int, input().split())
        u -= 1
        v -= 1
        g[u].append(v)
        g[v].append(u)
        d[u] += 1
        d[v] += 1
    q = [i for i in range(n) if d[i] == 1]
    dp = [1] * n
    ans = 1
    for v in q:
        d[v] -= 1
        for u in g[v]:
            if d[u] == 0:
                ans = (ans + dp[u]) % mod
                continue
            d[u] -= 1
            if d[u] == 1:
                q.append(u)
            dp[u] = dp[u] * (dp[v] + 1) % mod
    ans = (ans + dp[q[-1]]) % mod
    print(ans)

Solution (Python):
from collections import deque
R = lambda : map(int, input().split())
(t,) = R()

def answer():
    (n,) = R()
    adj_matrix = [[] for i in range(n + 1)]
    for _ in range(n - 1):
        (a, b) = R()
        adj_matrix[a].append(b)
        adj_matrix[b].append(a)
    stack = deque()
    stack.append((1, -1))
    route = []
    parents = [-1 for _ in range(n + 1)]
    while stack:
        (s, p) = stack.popleft()
        route.append(s)
        for v in adj_matrix[s]:
            if v == p:
                continue
            parents[v] = s
            stack.append((v, s))
    assert len(route) == n
    ans = [(1, 0) for _ in range(n + 1)]
    for v in route[::-1]:
        if len(adj_matrix[v]) == 1 and adj_matrix[v][0] == parents[v]:
            continue
        f1 = 0
        sf2 = 0
        sf1 = 1
        for u in adj_matrix[v]:
            if u == parents[v]:
                continue
            (vf1, vf2) = ans[u]
            f1 += vf1
            sf1 *= 1 + vf1
            sf1 = sf1 % 998244353
            sf2 += vf2
            sf2 = sf2 % 998244353
        sf1 -= f1
        sf1 -= 1
        rf1 = (1 + f1 + sf1) % 998244353
        rf2 = (sf2 + f1) % 998244353
        ans[v] = (rf1, rf2)
    (f1, f2) = ans[1]
    return (f1 + f2 + 1) % 998244353
for i in range(t):
    print(answer())

Solution (Python):
mod = 998244353
import sys
input = sys.stdin.readline
for _ in range(int(input())):
    n = int(input())
    g = [[] for k in range(n)]
    d = [0] * n
    for i in range(n - 1):
        (u, v) = map(int, input().split())
        u -= 1
        v -= 1
        g[u].append(v)
        g[v].append(u)
        d[u] += 1
        d[v] += 1
    q = [i for i in range(n) if d[i] == 1]
    dp = [1] * n
    ans = 1
    for v in q:
        d[v] -= 1
        for u in g[v]:
            if d[u] == 0:
                ans = (ans + dp[u]) % mod
                continue
            d[u] -= 1
            if d[u] == 1:
                q.append(u)
            dp[u] = dp[u] * (dp[v] + 1) % mod
    ans = (ans + dp[q[-1]]) % mod
    print(ans)