結果
問題 | 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 |
ソースコード
import syssys.setrecursionlimit(int(1e9))from typing import Callable, Generic, List, TypeVarT = 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 = nself._decrement = decrementdef addEdge(self, u: int, v: int) -> None:u -= self._decrementv -= self._decrementself.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._decrementassert 0 <= root < self._nparents = [-1] * self._norder = [root]stack = [root]while stack:cur = stack.pop()for next in self.adjList[cur]:if next == parents[cur]:continueparents[next] = curorder.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:continuedp2[next] = resres = op(res, composition(dp1[next], cur, next, 0))res = e(cur)for next in self.adjList[cur][::-1]:if parents[cur] == next:continuedp2[next] = op(res, dp2[next])res = op(res, composition(dp1[next], cur, next, 0))dp1[cur] = resfor 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 dp1def e(root: int) -> int:return 0def op(childRes1: int, childRes2: int) -> int:return childRes1 + childRes2def composition(fromRes: int, parent: int, cur: int, direction: int) -> int:if direction == 0: # cur -> parentreturn fromRes + subSize[cur]return fromRes + (n - subSize[cur]) # parent -> curdef dfsForSubSize(cur: int, parent: int) -> int:res = 1for next in R.adjList[cur]:if next != parent:res += dfsForSubSize(next, cur)subSize[cur] = resreturn resn = int(input())R = Rerooting(n)for _ in range(n - 1):u, v = map(int, input().split())R.addEdge(u - 1, v - 1)subSize = [0] * ndfsForSubSize(0, -1)dp = R.rerooting(e=e, op=op, composition=composition, root=0)print(sum(dp) + n * n)