gmo61520

Untitled

Apr 28th, 2016
173
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.18 KB | None | 0 0
  1. from collections import defaultdict, namedtuple
  2. Node = namedtuple('Node', 'id key left_child_id right_child_id color')
  3.  
  4. # parsing input
  5. nodes_count, root_node_id = map(int, input().split())
  6. tree = {}
  7. for n in range(nodes_count):
  8. node = Node(*map(int, input().split()))
  9. tree[node.id] = node
  10.  
  11. # building depths
  12. def build_depth(node_id, tree, storage, current_depth = 0):
  13. node = tree[node_id]
  14. storage[current_depth].append(node)
  15.  
  16. if node.left_child_id:
  17. build_depth(node.left_child_id, tree, storage, current_depth + 1)
  18. if node.right_child_id:
  19. build_depth(node.right_child_id, tree, storage, current_depth + 1)
  20.  
  21. nodes_by_depth = defaultdict(list)
  22. build_depth(root_node_id, tree, nodes_by_depth)
  23.  
  24. # calculating
  25. neighbours_count = 0
  26. for nodes in nodes_by_depth.values():
  27. last_black = Node(0, None, 0, 0, 1)
  28. last_white = Node(0, None, 0, 0, 0)
  29. for node in nodes:
  30. if node.color == 1:
  31. last_node, last_black = last_black, node
  32. else:
  33. last_node, last_white = last_white, node
  34. if last_node.key == node.key:
  35. neighbours_count += 1
  36.  
  37. # output result
  38. print(neighbour_pairs)
Advertisement
Add Comment
Please, Sign In to add comment