freee Programming Contest 2023（AtCoder Beginner Contest 310）G問題メモ
G - Takahashi And Pass-The-Ball Game
問題
$N$ 人の人がいて、はじめ、それぞれ $B_1,...,B_N$ 個のボールを持っている
数列 $A=(A_1,...,A_N)$ に従って、「人 $i$ は $A_i$ に自分が持っている全てのボールを渡す」という操作を1回とする
この操作を繰り返す回数を $1～K$ から一様ランダムに決める
全ての操作後、各人が持っているボールの個数の期待値を $\mod{998244353}$ で求めよ
$1 \le N \le 2 \times 10^5$
$1 \le K \le 10^{18}$
解法
求めるのは期待値だが、$K$ で割る前の値である「$1,2,...,K$ 回操作後に渡されるボールの個数の合計」が各人についてわかればよい。
つまり、各頂点 $i$ から $1,...,K$ 回後に訪れる頂点のそれぞれに $B_i$ を加算したい。
もちろんこれを愚直に行うことは時間制約上できない。
$i→A_i$ に辺を張ったグラフはFunctional Graphなので、各連結成分はループが1つとそこに流入する木からなることを利用する。
「ループ上」「木の根に向かうパス上」それぞれにおける区間加算は、累積和を用いて高速化できる。
加算したい頂点がループ上か木上かを区別して、別々に扱うことで実現できる。
ループ上の累積和
ループ上の累積和は、適当な点で切り開いて、「どこから累積和を始めるか」を固定することが重要となる。
その上で、「①最初の周回の末尾まで」「②周回分」「③残り」の3つに分けて加算するとやりやすい。
○→○→●→○→○
Bi
＋ ---- － ①
＋ ---------------- － x周回分 ②
＋ ---- － ③
木上の累積和
ループ上の各頂点を根とした木をDFSで辿る。
ループ r
→○→○→◎→○→○→
↑
◎
↗↖ i
○ ◎←◎←●
根 $r$ までのパスを持ちながらDFSする。木上の累積和で加算するのは
このようにし、葉から根に向かって累積和を取っていく（$Ans[A_i]+=Ans[i]$）と、
パスへの区間加算を表現できる。
Python3
from itertools import accumulate
def add_loop(m, loop_ans, si, k, b):
""" ループ内位置 si から k 個に渡って b を加算することを表現するように累積和を更新 """
if k <= 0:
return
loop_rem = m - si # 1周目の末端までの残り個数
if k <= loop_rem:
loop_ans[si] += b
loop_ans[si + k] -= b
return
loop_ans[si] += b
loop_ans[m] -= b
entire_loop, last_loop = divmod(k - loop_rem, m)
loop_ans[0] += b * (entire_loop + 1)
loop_ans[m] -= b * entire_loop
loop_ans[last_loop] -= b
def solve(loop):
m = len(loop)
loop_ans = [0] * (m + 1) # ループ内累積和
# 各ループ内頂点に付き、以下を累積和で更新
# ①自身からループ内への影響
# ②自身へ至るループ外の枝からループ内への影響
# ③自身へ至るループ外の枝から、同じ枝への影響
for li in range(m):
u = loop[li]
nxt_i = (li + 1) % m
prev_in_loop = loop[(li - 1) % m]
# ①
add_loop(m, loop_ans, nxt_i, k, bbb[loop[li]])
# ②③
path = [u]
while path:
v = path[-1]
if progress[v] == len(prev[v]):
if v != u:
# 自身の影響を反映
b = bbb[v]
ans[aaa[v]] += b
if len(path) > k + 1:
# ③ Kが同じ枝で終わる
ans[path[-k - 2]] -= b
else:
# ② Kがループ内に影響する
add_loop(m, loop_ans, nxt_i, k - len(path) + 1, b)
# 次へ累積和を引き継ぐ
ans[v] %= MOD
ans[aaa[v]] += ans[v]
path.pop()
else:
w = prev[v][progress[v]]
progress[v] += 1
if w == prev_in_loop:
continue
path.append(w)
loop_ans = list(accumulate(loop_ans))
for li, v in enumerate(loop):
ans[v] += loop_ans[li]
ans[v] %= MOD
n, k = map(int, input().split())
aaa = list(map(int, input().split()))
aaa = [a - 1 for a in aaa]
bbb = list(map(int, input().split()))
MOD = 998244353
prev = [[] for _ in range(n)]
for i, a in enumerate(aaa):
prev[a].append(i)
checked = [False] * n
progress = [0] * n
ans = [0] * n
for s in range(n):
if checked[s]:
continue
# ループを検出
path = []
v = s
while checked[v] == False:
checked[v] = True
path.append(v)
v = aaa[v]
try:
loop_start = path.index(v)
except:
# 到達できるループは既に他の枝から調べられているので、調べる必要なし
continue
loop = path[loop_start:]
solve(loop)
inv_k = pow(k, MOD - 2, MOD)
ans = [a * inv_k % MOD for a in ans]
print(*ans)
