# 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)) ```