結果

問題 No.2949 Product on Tree
コンテスト
ユーザー koheijkt
提出日時 2026-01-13 21:39:23
言語 PyPy3
(7.3.17)
結果
AC  
実行時間 1,735 ms / 2,000 ms
コード長 1,457 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 336 ms
コンパイル使用メモリ 82,704 KB
実行使用メモリ 335,820 KB
最終ジャッジ日時 2026-01-13 21:40:23
合計ジャッジ時間 53,306 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 46
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

import sys
sys.setrecursionlimit(10**8)
import pypyjit
pypyjit.set_param('max_unroll_recursion=-1')
MOD = 998244353
N = int(input())
A = [0] + list(map(int, input().split()))
G = [list() for _ in range(N + 1)]
for i in range(N - 1):
    u, v = map(int, input().split())
    G[u].append(v)
    G[v].append(u)

# DP[i]
# 部分木iにおいて、f(i, 部分木i内の全頂点) の合計
DP = [0] * (N + 1)
def dfs(pos, pre):
    for nex in G[pos]:
        if nex == pre:
            continue
        dfs(nex, pos)
        # 親に計上
        DP[pos] += DP[nex]
    # 子どもの計上がすべて完了したら、+1 してから A[i] を掛ける
    DP[pos] = (DP[pos] + 1) * A[pos] % MOD
    return

dfs(1, 0)

# 全方位木DPを解く
DP2 = [0] * (N + 1)
def dfs2(pos, pre):
    # pos が根になったときの、DP が答え
    DP2[pos] = DP[pos]
    # 根の移動
    for nex in G[pos]:
        if nex == pre:
            continue
        # 設定のバックアップ
        memo1, memo2 = DP[pos], DP[nex]
        # DP[pos] からは DP[nex] 分の寄与を削る
        DP[pos] -= A[pos] * DP[nex]
        DP[pos] %= MOD
        # DP[nex] には、DP[pos]（更新後）の寄与を追加する
        DP[nex] += A[nex] * DP[pos]
        DP[nex] %= MOD
        dfs2(nex, pos)
        # リストアする
        DP[pos], DP[nex] = memo1, memo2
    return

dfs2(1, 0)
ans = sum(DP2) - sum(A)
ans *= pow(2, -1, MOD)
ans %= MOD
print(ans)
0