結果
問題 | No.1333 Squared Sum |
ユーザー |
![]() |
提出日時 | 2021-01-08 22:54:05 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 1,149 ms / 2,000 ms |
コード長 | 1,580 bytes |
コンパイル時間 | 335 ms |
コンパイル使用メモリ | 82,432 KB |
実行使用メモリ | 197,152 KB |
最終ジャッジ日時 | 2024-11-16 14:37:13 |
合計ジャッジ時間 | 31,060 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge4 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
other | AC * 44 |
ソースコード
# 全方位木DP import sys input = sys.stdin.readline from collections import deque mod = 10 ** 9 + 7 N = int(input()) X = [[] for i in range(N)] for i in range(N-1): x, y, w = map(int, input().split()) X[x-1].append((y-1, w)) X[y-1].append((x-1, w)) P = [-1] * N Q = deque([0]) R = [] D = [0] * N while Q: i = deque.popleft(Q) R.append(i) for a, w in X[i]: if a != P[i]: P[a] = i D[a] = w X[a].remove((i, w)) deque.append(Q, a) # # a0 += b0 # a1 += b1 + b0 * d # a2 += b0 * d ** 2 + 2 * b1 * d + b2 ##### Settings unit = (0, 0, 0) merge = lambda a, b: ((a[0] + b[0]) % mod, (a[1] + b[1]) % mod, (a[2] + b[2]) % mod) ## ↓ この2つは必ず対称的にする(じゃないとバグる) def adj_bu(a, i, p): a0, a1, a2 = a a0 += 1 d = D[i] return (a0, (a1 + a0 * d) % mod, (a0 * d ** 2 + 2 * a1 * d + a2) % mod) def adj_td(a, i, p): a0, a1, a2 = a a0 += 1 d = D[i] return (a0, (a1 + a0 * d) % mod, (a0 * d ** 2 + 2 * a1 * d + a2) % mod) adj_fin = lambda a, i: a[2] ##### ME = [unit] * N XX = [0] * N TD = [unit] * N for i in R[1:][::-1]: p = P[i] XX[i] = adj_bu(ME[i], i, p) ME[p] = merge(ME[p], XX[i]) XX[R[0]] = adj_fin(ME[R[0]], R[0]) for i in R: ac = TD[i] for j, _ in X[i]: TD[j] = ac ac = merge(ac, XX[j]) ac = unit for j, _ in X[i][::-1]: TD[j] = adj_td(merge(TD[j], ac), j, i) ac = merge(ac, XX[j]) XX[j] = adj_fin(merge(ME[j], TD[j]), j) print(sum(XX) * (mod + 1) // 2 % mod)