結果

問題 No.1418 Sum of Sum of Subtree Size
ユーザー 草苺奶昔草苺奶昔
提出日時 2023-03-13 19:39:27
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 497 ms / 2,000 ms
コード長 2,802 bytes
コンパイル時間 231 ms
コンパイル使用メモリ 82,344 KB
実行使用メモリ 178,816 KB
最終ジャッジ日時 2024-09-18 07:41:04
合計ジャッジ時間 11,159 ms
ジャッジサーバーID
(参考情報)
judge1 / judge5
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 41
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

import sys
sys.setrecursionlimit(int(1e9))
from typing import Callable, Generic, List, TypeVar
T = TypeVar("T")
class Rerooting(Generic[T]):
__slots__ = ("adjList", "_n", "_decrement")
def __init__(self, n: int, decrement: int = 0):
self.adjList = [[] for _ in range(n)]
self._n = n
self._decrement = decrement
def addEdge(self, u: int, v: int) -> None:
u -= self._decrement
v -= self._decrement
self.adjList[u].append(v)
self.adjList[v].append(u)
def rerooting(
self,
e: Callable[[int], T],
op: Callable[[T, T], T],
composition: Callable[[T, int, int, int], T],
root=0,
) -> List["T"]:
root -= self._decrement
assert 0 <= root < self._n
parents = [-1] * self._n
order = [root]
stack = [root]
while stack:
cur = stack.pop()
for next in self.adjList[cur]:
if next == parents[cur]:
continue
parents[next] = cur
order.append(next)
stack.append(next)
dp1 = [e(i) for i in range(self._n)]
dp2 = [e(i) for i in range(self._n)]
for cur in order[::-1]:
res = e(cur)
for next in self.adjList[cur]:
if parents[cur] == next:
continue
dp2[next] = res
res = op(res, composition(dp1[next], cur, next, 0))
res = e(cur)
for next in self.adjList[cur][::-1]:
if parents[cur] == next:
continue
dp2[next] = op(res, dp2[next])
res = op(res, composition(dp1[next], cur, next, 0))
dp1[cur] = res
for newRoot in order[1:]:
parent = parents[newRoot]
dp2[newRoot] = composition(op(dp2[newRoot], dp2[parent]), parent, newRoot, 1)
dp1[newRoot] = op(dp1[newRoot], dp2[newRoot])
return dp1
def e(root: int) -> int:
return 0
def op(childRes1: int, childRes2: int) -> int:
return childRes1 + childRes2
def composition(fromRes: int, parent: int, cur: int, direction: int) -> int:
if direction == 0: # cur -> parent
return fromRes + subSize[cur]
return fromRes + (n - subSize[cur]) # parent -> cur
def dfsForSubSize(cur: int, parent: int) -> int:
res = 1
for next in R.adjList[cur]:
if next != parent:
res += dfsForSubSize(next, cur)
subSize[cur] = res
return res
n = int(input())
R = Rerooting(n)
for _ in range(n - 1):
u, v = map(int, input().split())
R.addEdge(u - 1, v - 1)
subSize = [0] * n
dfsForSubSize(0, -1)
dp = R.rerooting(e=e, op=op, composition=composition, root=0)
print(sum(dp) + n * n)
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0