結果

問題 No.2949 Product on Tree
ユーザー hato336hato336
提出日時 2024-10-25 21:53:57
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 958 ms / 2,000 ms
コード長 2,488 bytes
コンパイル時間 544 ms
コンパイル使用メモリ 82,664 KB
実行使用メモリ 150,056 KB
最終ジャッジ日時 2024-10-25 21:54:43
合計ジャッジ時間 42,747 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 46
権限があれば一括ダウンロードができます

ソースコード

diff # 

import collections,sys,math,functools,operator,itertools,bisect,heapq,decimal,string,time,random
#sys.setrecursionlimit(10**9)
#sys.set_int_max_str_digits(0)
#input = sys.stdin.readline

#alist = list(map(int,input().split()))
class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        group_members = collections.defaultdict(list)
        for member in range(self.n):
            group_members[self.find(member)].append(member)
        return group_members

    def __str__(self):
        return ''.join(f'{r}: {m}' for r, m in self.all_group_members().items())
#alist = []
#s = input()
n = int(input())
a = list(map(int,input().split()))
edge = [[] for i in range(n)]
for i in range(n-1):
    u,v = map(int,input().split())
    u-=1
    v-=1
    edge[u].append(v)
    edge[v].append(u)
mod = 998244353
#https://qiita.com/Kiri8128/items/cbaa021dbcb07b5fdb92#非再帰-bfs
plist = [-1 for i in range(n)]
d = collections.deque()
r = []
d.append(0)
while d:
    now = d.popleft()
    r.append(now)
    for i in edge[now]:
        plist[i] = now
        d.append(i)
        edge[i].remove(now)
dp = [0 for i in range(n)]
ans = 0
for i in r[::-1]:
    b = []
    for j in edge[i]:
        dp[i] += dp[j] * a[i]
        dp[i] %= mod
        ans += dp[j] * a[i]
        ans %= mod

        b.append(dp[j])
    dp[i] += a[i]
    if not b:
        continue
    c = functools.reduce(lambda x,y:(x+y)%mod,b)
    c **= 2
    c %= mod
    for j in b:
        c -= j**2
        c %= mod
    c *= pow(2,mod-2,mod)
    c %= mod
    ans += c * a[i]
    ans %= mod

print(ans)
0