# Colleage and Supervisor relations
```python=
from collections import defaultdict
class Node:
def __init__(self, val):
self.val = val
self.childs = set()
self.parent = None
def build_tree(relations):
all_nodes = {}
for r_type, A, B in relations:
if A not in all_nodes:
all_nodes[A] = Node(A)
node_a = all_nodes[A]
if B not in all_nodes:
all_nodes[B] = Node(B)
node_b = all_nodes[B]
if r_type == 0: # A, B are colleages
node_a_parent = node_a.parent
node_b_parent = node_b.parent
if node_a_parent is None and node_b_parent is None:
parent = Node(None)
else:
if node_a_parent:
parent = node_a_parent
else:
parent = node_b_parent
node_a.parent = parent
node_b.parent = parent
parent.childs.add(node_a)
parent.childs.add(node_b)
else: # A is supervisor of B
if node_b.parent is None:
node_b.parent = node_a
else:
node_a = node_b.parent
node_a.val = A
all_nodes[A] = node_a
node_a.childs.add(node_b)
return all_nodes
def build_tree_table(all_nodes):
# [root1, root2, ...]
tree_table = []
# node_infos[$node] = [$tree_number, $in, $out]
node_infos = defaultdict(lambda: [-1, -1, -1])
for node_name in all_nodes:
node = all_nodes[node_name]
if node in node_infos:
continue
vis = set()
table_number = -1
curr = node
while curr:
if curr in node_infos:
table_number = node_infos[curr][0]
break
vis.add(curr)
root = curr
curr = curr.parent
if table_number == -1:
table_number = len(tree_table)
tree_table.append(root)
for vis_node in vis:
node_infos[vis_node][0] = table_number
return tree_table, node_infos
def build_relations(tree_table, node_infos):
def dfs(node):
if not node:
return
nonlocal cnt
cnt += 1
node_infos[node][1] = cnt # in
for next_node in node.childs:
dfs(next_node)
cnt += 1
node_infos[node][2] = cnt # out
for root in tree_table:
cnt = 0
dfs(root)
def solve(relations, Q):
# N: node's number
# M: Queries length
# O(N)
all_nodes = build_tree(relations)
# O(N)
tree_table, node_infos = build_tree_table(all_nodes)
# O(N)
build_relations(tree_table, node_infos)
# O(M)
ans = []
for q_type, A, B in Q:
if A not in all_nodes or B not in all_nodes:
ans.append(False)
continue
node_a, node_b = all_nodes[A], all_nodes[B]
# A, B are colleage?
if q_type == 0:
if node_a is not None and node_b is not None\
and node_a.parent == node_b.parent:
ans.append(True)
else:
ans.append(False)
# A is B's supervisor?
else:
if node_infos[node_a][0] == node_infos[node_b][0]:
if node_infos[node_a][1] < node_infos[node_b][1] and\
node_infos[node_a][2] > node_infos[node_b][2]:
ans.append(True)
else:
ans.append(False)
else:
ans.append(False)
return ans
relations = [
# Tree 1
# D
# ├── A
# ├── B
# └── C
[0, "A", "B"], # colleage
[0, "B", "C"], # colleage
[1, "D", "A"], # D is A's supervisor
# Tree 2
# D2
# ├── A2
# │ └── E2
# ├── B2
# └── C2
[0, "A2", "B2"], # colleage
[0, "B2", "C2"], # colleage
[1, "D2", "A2"], # D2 is A2's supervisor
[1, "A2", "E2"], # A2 is E2's supervisor
]
Q = [
[0, "A", "B"], # A is colleage of B?
[1, "D", "B"], # D is supervisor of B?
[1, "F", "B"], # F is supervisor of B?
[1, "D2", "B"], # D2 is supervisor of B?
[1, "D2", "E2"], # D2 is supervisor of E2?
]
print(solve(relations, Q))
```