結果
問題 | No.1103 Directed Length Sum |
ユーザー |
![]() |
提出日時 | 2022-11-20 12:55:19 |
言語 | PyPy3 (7.3.15) |
結果 |
TLE
|
実行時間 | - |
コード長 | 3,620 bytes |
コンパイル時間 | 249 ms |
コンパイル使用メモリ | 81,692 KB |
実行使用メモリ | 410,468 KB |
最終ジャッジ日時 | 2024-09-21 13:27:11 |
合計ジャッジ時間 | 14,684 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge2 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 2 |
other | AC * 21 TLE * 1 |
ソースコード
from typing import List, Tuple, Callable, TypeVarfrom typing import List, Tuple, Optionalimport sysimport itertoolsimport heapqimport bisectimport mathfrom collections import deque, defaultdictfrom functools import lru_cache, cmp_to_keyinput = sys.stdin.readlineif __file__ != 'prog.py':sys.setrecursionlimit(10 ** 6)def readints(): return map(int, input().split())def readlist(): return list(readints())def readstr(): return input()[:-1]T = TypeVar('T')input = sys.stdin.readlineclass Rerooting:# reference: https://null-mn.hatenablog.com/entry/2020/04/14/124151# 適当な頂点vを根とする部分木に対して計算される値dp_vが、vの子c1, c2, ... ckを用いて# 下記のように表すことができる# dp_v = g(merge(f(dp_c1,c1), f(dp_c2,c2), ..., f(dp_ck,ck)), v)def __init__(self, N: int, E: List[Tuple[int, int]],f: Callable[[T, int, int, int], T],g: Callable[[T, int], T],merge: Callable[[T, T], T], e: T,root: int = 0):self.N = Nself.E = Eself.f = fself.g = gself.merge = mergeself.e = eself.dp = [[self.e for _ in range(len(self.E[v]))] for v in range(self.N)]self._calculate(root)def _dfs1(self, root):stack = [(root, -1)]ret = [self.e] * self.Nwhile stack:v, p = stack.pop()if v < 0:v = ~vacc = self.efor i, (c, d) in enumerate(self.E[v]):if d == p:continueself.dp[v][i] = ret[d]acc = self.merge(acc, self.f(ret[d], v, d, c))ret[v] = self.g(acc, v)continuestack.append((~v, p))for i, (c, d) in enumerate(self.E[v]):if d == p:continuestack.append((d, v))def _dfs2(self, root):stack = [(root, -1, self.e)]while stack:v, p, from_par = stack.pop()for i, (c, d) in enumerate(self.E[v]):if d == p:self.dp[v][i] = from_parbreakch = len(self.E[v])Sr = [self.e] * (ch + 1)for i in range(ch, 0, -1):c, d = self.E[v][i - 1]Sr[i - 1] = self.merge(Sr[i], self.f(self.dp[v][i - 1], v, d, c))Sl = self.efor i, (c, d) in enumerate(self.E[v]):if d != p:val = self.merge(Sl, Sr[i + 1])stack.append((d, v, self.g(val, v)))Sl = self.merge(Sl, self.f(self.dp[v][i], v, d, c))def _calculate(self, root):self._dfs1(root)self._dfs2(root)def solve(self, v):ans = self.efor i, (c, d) in enumerate(self.E[v]):ans = self.merge(ans, self.f(self.dp[v][i], v, d, c))return self.g(ans, v)N = int(input())E = [[] for _ in range(N)]D = [0] * Nfor _ in range(N - 1):a, b = readints()a -= 1b -= 1E[a].append((1, b))D[b] += 1for i in range(N):if D[i] == 0:root = ibreakmod = 10 ** 9 + 7def f(a, v, ch, cost):return ((a[0] + a[1]) % mod, a[1])def g(a, v):return (a[0], a[1] + 1)def merge(a, b):return ((a[0] + b[0]) % mod, a[1] + b[1])solver = Rerooting(N, E, f, g, merge, (0, 0), root=root)ans = 0for i in range(N):ans += solver.solve(i)[0]ans %= modprint(ans)