結果

問題 No.2949 Product on Tree
コンテスト
ユーザー koheijkt
提出日時 2026-01-13 21:43:07
言語 Python3
(3.14.2 + numpy 2.4.0 + scipy 1.16.3)
結果
AC  
実行時間 1,966 ms / 2,000 ms
コード長 1,459 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 592 ms
コンパイル使用メモリ 20,916 KB
実行使用メモリ 102,388 KB
最終ジャッジ日時 2026-01-13 21:44:38
合計ジャッジ時間 82,885 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
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