Problem F — Random Walk — Codeforces Round #868

Revision en3, by rhymehatch, 2024-03-15 05:09:41

Editorial for this problem and solutions in the comments rely on a lot of symbolic expansion of various expressions until they discover regularity in formulas.

I'm used to looking at random walks as recurrences and directly solving them.

Tree below formulates the system of equations below where $$$p_i$$$ is the expectation of counts:

$$$ \begin{align*} p_1 &= \frac{p_3}{2} \\ p_2 &= 1 + \frac{p_3}{2} + p_8 \\ p_3 &= p_1 + \frac{p_2}{3} \\ p_4 &= p_5 + p_7 \\ p_5 &= \frac{p_4}{3} \\ p_6 &= 1 \\ p_7 &= \frac{p_4}{3} \\ p_8 &= \frac{p_2}{3} \end{align*} $$$

Solving general system of equations with $$$n$$$ variables is slow. I almost quit trying but this is a tree, and the rabbit out of the hat is that this sparse system can be solved by tree traversal. We start by pushing the coefficients at leaves to the parents and eventually get to the starting node (in tree above it is 2), and we can calculate the coefficient. Then a single BFS can push this information back to all nodes and calculate their expectation.

This is a general method to solve any system of equations defined by a tree.

Division is handled through pow(x, -1, M).

n, s, t, *e = map(int, open(0).read().split())
G = [[] for _ in -~n*' ']
for a, b in zip(*[iter(e)]*2):
	G[a] += b,
	G[b] += a,

M = 998244353

V = *C, = -~n*[0]

i = lambda x: pow(x, -1, M)  # handling division as multiplication
Q = [(s, 0)]  # only way to do stackless dfs in Python
while Q:
	x, p = Q[-1]
	if V[x]<1:
		V[x] = 1
		for u in G[x]:
			if u!=t and u!=p: Q += (u, x),
	else:
		c = 0
		for u in G[x]:
			if u!=p and u!=t: c = (c+C[u]*i(len(G[u])))%M  # p_i = 1/degree(child) * p_child
		C[x] = i([len(G[p]), 1][p==0])*i(M+1-c)%M  # (1 - c) * p_i = 1/degree(parent) * p_{parent}
                # final coefficient C[x] is 1/((1-c)*degree(parent)) so when C[p] is materialized, multiply
		Q.pop()

# C[s] is the only coefficient that is currently calculated
B = [(s, 0)]  # short BFS
for x, p in B:
	for u in G[x]:
		if u!=t and u!=p:
			C[u] = C[u]*C[x]%M
			B += (u, x),
C[t] = 1
print(*C[1:])

P.S. Python 3.11 comes with 100000x efficient stack (cframes improvement). It will be possible to write a much shorter implementation with recursive DFS.

Tags equations, bfs, dfs, expectation

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en8 English rhymehatch 2024-03-19 15:27:42 34 fix wrong indentation
en7 English rhymehatch 2024-03-19 15:24:46 49 added link to problem statement
en6 English rhymehatch 2024-03-18 17:12:36 53 simpler C[x] calculation using 1 modular inverse calculation + tags
en5 English rhymehatch 2024-03-15 11:58:56 33
en4 English rhymehatch 2024-03-15 11:36:12 2
en3 English rhymehatch 2024-03-15 05:09:41 157
en2 English rhymehatch 2024-03-15 05:06:38 11
en1 English rhymehatch 2024-03-15 05:05:51 2248 Initial revision (published)