結果
問題 | No.2892 Lime and Karin |
ユーザー | adapchi |
提出日時 | 2024-09-13 23:31:20 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 3,936 ms / 8,000 ms |
コード長 | 1,395 bytes |
コンパイル時間 | 360 ms |
コンパイル使用メモリ | 82,304 KB |
実行使用メモリ | 144,400 KB |
最終ジャッジ日時 | 2024-09-13 23:33:00 |
合計ジャッジ時間 | 96,227 ms |
ジャッジサーバーID (参考情報) |
judge1 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 52 |
ソースコード
N = int(input())G = [[] for _ in range(N)]for _ in range(N - 1):u, v = [int(x) - 1 for x in input().split()]G[u].append(v)G[v].append(u)S = [1 if x == '1' else -1 for x in input()]used = [False] * Nsize = [0] * Ndef centroid(x):stk = [(x, None, 0)]while stk:u, p, t = stk.pop()if t == 0:stk.append((u, p, 1))stk.extend((v, u, 0) for v in G[u] if not used[v] and v != p)else:size[u] = 1 + sum(size[v] for v in G[u] if not used[v] and v != p)p = Nones = size[x]while True:y = next((y for y in G[x] if not used[y] and y != p and size[y] * 2 > s), None)if y is None: breakx, p = y, xreturn xdef paths(x, p):stk = [(x, p, S[x])]res = []while stk:u, p, d = stk.pop()res.append(d)stk.extend((v, u, d + S[v]) for v in G[u] if not used[v] and v != p)return resdef pairs(xs, y):xs.sort()res = 0i = len(xs)for x in xs:while i >= 1 and x + xs[i - 1] + y > 0:i -= 1res += len(xs) - ireturn resstk = [0]ans = 0while stk:u = centroid(stk.pop())if S[u] == 1: ans += 1vs = [v for v in G[u] if not used[v]]stk.extend(vs)ds = [paths(v, u) for v in vs]ans += sum(sum(x + S[u] > 0 for x in d) for d in ds)ans += (pairs(sum(ds, start=[]), S[u]) - sum(pairs(d, S[u]) for d in ds)) // 2used[u] = Trueprint(ans)