-
Notifications
You must be signed in to change notification settings - Fork 5
/
chainblock.py
98 lines (88 loc) · 3.26 KB
/
chainblock.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
# Copyright (c) 2021 kamyu. All rights reserved.
#
# Facebook Hacker Cup 2021 Round 2 - Problem B. Chainblock
# https://www.facebook.com/codingcompetitions/hacker-cup/2021/round-2/problems/B
#
# Time: O(NlogN)
# Space: O(N)
#
from functools import partial
# Template:
# https://github.com/kamyu104/GoogleCodeJam-2020/blob/master/Round%202/emacs++2_concise.py
class TreeInfos(object): # Time: O(NlogN), Space: O(NlogN), N is the number of nodes
def __init__(self, children):
def preprocess(curr, parent):
# visited order of the nodes
if curr:
O.append(curr) # modified
# depth of the node i
D[curr] = 1 if parent == -1 else D[parent]+1
# ancestors of the node i
if parent != -1:
P[curr].append(parent)
i = 0
while i < len(P[curr]) and i < len(P[P[curr][i]]):
P[curr].append(P[P[curr][i]][i])
i += 1
# the subtree of the node i is represented by traversal index L[i]..R[i]
C[0] += 1
L[curr] = C[0]
def divide(curr, parent):
stk.append(partial(postprocess, curr))
for i in reversed(xrange(len(children[curr]))):
child = children[curr][i]
if child == parent:
continue
stk.append(partial(divide, child, curr))
stk.append(partial(preprocess, curr, parent))
def postprocess(curr):
R[curr] = C[0]
N = len(children)
L, R, D, P, C, O = [0]*N, [0]*N, [0]*N, [[] for _ in xrange(N)], [-1], []
stk = []
stk.append(partial(divide, 0, -1))
while stk:
stk.pop()()
assert(C[0] == N-1)
self.L, self.R, self.D, self.P, self.O = L, R, D, P, O
# Template:
# https://github.com/kamyu104/FacebookHackerCup-2019/blob/master/Final%20Round/little_boat_on_the_sea.py
def is_ancestor(self, a, b): # includes itself
return self.L[a] <= self.L[b] <= self.R[b] <= self.R[a]
def lca(self, a, b):
if self.D[a] > self.D[b]:
a, b = b, a
if self.is_ancestor(a, b):
return a
for i in reversed(xrange(len(self.P[a]))): # O(logN)
if i < len(self.P[a]) and not self.is_ancestor(self.P[a][i], b):
a = self.P[a][i]
return self.P[a][0]
def chainblock():
N = input()
adj = [[] for _ in xrange(N)]
for _ in xrange(N-1):
A, B = map(int, raw_input().strip().split())
adj[A-1].append(B-1)
adj[B-1].append(A-1)
F = map(lambda x: int(x)-1, raw_input().strip().split())
tree_infos = TreeInfos(adj)
groups = [[] for _ in xrange(N)]
for i, f in enumerate(F):
groups[f].append(i)
A = [float("inf")]*N
for f, idxs in enumerate(groups):
if not idxs:
continue
lca = reduce(lambda lca, x: tree_infos.lca(lca, x), idxs)
for i in idxs:
A[i] = tree_infos.D[lca]
result = 0
for i in reversed(tree_infos.O):
if A[i] == tree_infos.D[i]:
result += 1
else:
A[tree_infos.P[i][0]] = min(A[tree_infos.P[i][0]], A[i])
return result
for case in xrange(input()):
print 'Case #%d: %s' % (case+1, chainblock())