freee Programming Contest 2023（AtCoder Beginner Contest 310）G問題メモ

freee Programming Contest 2023（AtCoder Beginner Contest 310）

男子達が水泳をやる会社ではなく、会計管理サービスを中心に提供する会社。

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する。木上の累積和で加算するのは

$i$ のすぐ次の頂点 $A_i$ に $Ans[A_i]+=B_i$
$K$ 回目の操作で渡す対象が、まだ根を除く木上かどうかで変化（パスの長さで判断できる）
- 木上なら、その親頂点を $p$ とし、 $Ans[p]-=B_i$
- ループに入るなら、残り回数（ $K-$ 根までの距離）だけ、 $r$ を始点としてループ上の累積和に加算

このようにし、葉から根に向かって累積和を取っていく（ $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)

            

        
      