結果

問題 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
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 128 ms
85,876 KB
testcase_01 AC 129 ms
85,824 KB
testcase_02 AC 142 ms
89,352 KB
testcase_03 AC 908 ms
137,768 KB
testcase_04 AC 906 ms
132,872 KB
testcase_05 AC 891 ms
137,796 KB
testcase_06 AC 909 ms
137,948 KB
testcase_07 AC 841 ms
133,556 KB
testcase_08 AC 924 ms
137,720 KB
testcase_09 AC 876 ms
137,684 KB
testcase_10 AC 888 ms
137,524 KB
testcase_11 AC 907 ms
137,280 KB
testcase_12 AC 902 ms
135,784 KB
testcase_13 AC 903 ms
135,168 KB
testcase_14 AC 873 ms
129,660 KB
testcase_15 AC 883 ms
128,700 KB
testcase_16 AC 868 ms
130,272 KB
testcase_17 AC 878 ms
131,816 KB
testcase_18 AC 918 ms
133,876 KB
testcase_19 AC 878 ms
134,956 KB
testcase_20 AC 938 ms
135,804 KB
testcase_21 AC 905 ms
137,216 KB
testcase_22 AC 886 ms
130,360 KB
testcase_23 AC 944 ms
138,648 KB
testcase_24 AC 958 ms
138,260 KB
testcase_25 AC 924 ms
138,260 KB
testcase_26 AC 937 ms
138,128 KB
testcase_27 AC 912 ms
138,380 KB
testcase_28 AC 932 ms
138,472 KB
testcase_29 AC 906 ms
138,348 KB
testcase_30 AC 924 ms
138,132 KB
testcase_31 AC 914 ms
138,264 KB
testcase_32 AC 916 ms
138,128 KB
testcase_33 AC 924 ms
138,000 KB
testcase_34 AC 939 ms
137,996 KB
testcase_35 AC 902 ms
137,876 KB
testcase_36 AC 915 ms
137,868 KB
testcase_37 AC 909 ms
137,624 KB
testcase_38 AC 926 ms
137,620 KB
testcase_39 AC 940 ms
137,744 KB
testcase_40 AC 903 ms
137,624 KB
testcase_41 AC 938 ms
137,236 KB
testcase_42 AC 909 ms
137,496 KB
testcase_43 AC 353 ms
142,700 KB
testcase_44 AC 350 ms
141,292 KB
testcase_45 AC 390 ms
143,432 KB
testcase_46 AC 391 ms
137,904 KB
testcase_47 AC 322 ms
130,056 KB
testcase_48 AC 383 ms
150,056 KB
権限があれば一括ダウンロードができます

ソースコード

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