結果

問題 No.2004 Incremental Coins
ユーザー lam6er
提出日時 2025-03-26 15:50:24
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,393 bytes
コンパイル時間 328 ms
コンパイル使用メモリ 82,452 KB
実行使用メモリ 77,472 KB
最終ジャッジ日時 2025-03-26 15:51:15
合計ジャッジ時間 5,580 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 6 TLE * 1 -- * 13
権限があれば一括ダウンロードができます

ソースコード

diff # 

MOD = 998244353

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    ptr = 0
    N = int(data[ptr])
    ptr += 1
    K = int(data[ptr])
    ptr += 1
    
    A = list(map(int, data[ptr:ptr + N + 1]))
    ptr += N + 1
    
    P = list(map(int, data[ptr:ptr + N]))
    ptr += N
    
    # Compute depth for each node (0-based)
    depth = [0] * (N + 1)
    for j in range(1, N + 1):
        depth[j] = depth[P[j-1]] + 1
    
    # Binary Lifting (Doubling) setup for ancestors
    LOG = 20  # Enough for 2^20 > 2e5
    up = [[-1] * (N + 1) for _ in range(LOG)]
    # up[k][j] is 2^k ancestor of j
    # j ranges from 0 to N
    for j in range(N + 1):
        up[0][j] = P[j-1] if j != 0 else -1  # j=0 has no parent
    
    for k in range(1, LOG):
        for j in range(N + 1):
            if up[k-1][j] == -1:
                up[k][j] = -1
            else:
                up[k][j] = up[k-1][up[k-1][j]]
    
    # Precompute inverses
    max_L = max(depth)
    L_max = min(max_L, K)
    inv = [1] * (L_max + 2)
    for i in range(1, L_max + 1):
        inv[i] = pow(i, MOD - 2, MOD)
    
    # Precompute comb[L] = C(K, L) mod MOD
    comb = [0] * (L_max + 1)
    comb[0] = 1
    for L in range(1, L_max + 1):
        term = (K - L + 1) % MOD
        term = (term + MOD) % MOD  # Ensure non-negative
        comb_L = comb[L-1] * term % MOD
        comb_L = comb_L * inv[L] % MOD
        comb[L] = comb_L
    
    # Now, for each node j, add A[j] * comb[L] to each ancestor i at distance L
    ans = [0] * (N + 1)
    for j in range(N + 1):
        current = j
        max_L_j = min(depth[j], K)
        for L in range(0, max_L_j + 1):
            # Find ancestor at distance L from j
            # To find L-th ancestor of j
            if L == 0:
                i = j
            else:
                i = j
                cnt = L
                for k in reversed(range(LOG)):
                    if cnt >= (1 << k):
                        cnt -= (1 << k)
                        i = up[k][i]
                        if i == -1:
                            break
                if cnt != 0 or i == -1:
                    # No such ancestor
                    continue
            # Add A[j] * comb[L] to ans[i]
            ans[i] = (ans[i] + A[j] * comb[L]) % MOD
    
    for x in ans:
        print(x % MOD)

if __name__ == "__main__":
    main()
0