Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from collections import defaultdict, namedtuple
- Node = namedtuple('Node', 'id key left_child_id right_child_id color')
- # parsing input
- nodes_count, root_node_id = map(int, input().split())
- tree = {}
- for n in range(nodes_count):
- node = Node(*map(int, input().split()))
- tree[node.id] = node
- # building depths
- def build_depth(node_id, tree, storage, current_depth = 0):
- node = tree[node_id]
- storage[current_depth].append(node)
- if node.left_child_id:
- build_depth(node.left_child_id, tree, storage, current_depth + 1)
- if node.right_child_id:
- build_depth(node.right_child_id, tree, storage, current_depth + 1)
- nodes_by_depth = defaultdict(list)
- build_depth(root_node_id, tree, nodes_by_depth)
- # calculating
- neighbours_count = 0
- for nodes in nodes_by_depth.values():
- last_black = Node(0, None, 0, 0, 1)
- last_white = Node(0, None, 0, 0, 0)
- for node in nodes:
- if node.color == 1:
- last_node, last_black = last_black, node
- else:
- last_node, last_white = last_white, node
- if last_node.key == node.key:
- neighbours_count += 1
- # output result
- print(neighbours_count)
Advertisement
Add Comment
Please, Sign In to add comment