Skip to content

Instantly share code, notes, and snippets.

# aglee/TreeDiff.py Created Oct 6, 2016

 #!/usr/bin/python """ https://www.hackerrank.com/challenges/cut-the-tree """ class Node(object): def __init__(self, weight): self.weight = weight self.adjacent_nodes = [] self.subtree_weight = weight self.is_queued = False self.parent_node = None # First line contains number of nodes. node_count = int(raw_input()) # Next line contains the node weights. # Create an array of nodes with those weights. nodes = list(map(lambda x: Node(int(x)), raw_input().split())) # The next n-1 lines contain pairs of node indexes indicating edges. # Note that these indexes are 1-based. for _ in range(node_count - 1): edge_indexes = list(map(lambda x: int(x), raw_input().split())) node_one = nodes[edge_indexes - 1] node_two = nodes[edge_indexes - 1] node_one.adjacent_nodes.append(node_two) node_two.adjacent_nodes.append(node_one) # Arrange the nodes in breadth-first order. Here we designate nodes # as the root node of the entire tree, but since the edges of a tree are # non-directed (i.e. there is no inherent parent-child relationship), # we could use any other node equally well. root_node_index = 0 queue = [nodes[root_node_index]] nodes[root_node_index].is_queued = True queue_index = 0 while queue_index < len(queue): nd = queue[queue_index] for adj in nd.adjacent_nodes: if not adj.is_queued: queue.append(adj) adj.is_queued = True adj.parent_node = nd queue_index += 1 # Add every node's weight to its parent's subtree_weight. # Work up from the bottom of the tree. Note that every node's # subtree_weight was initialized to its own weight. for node in reversed(queue): if node.parent_node: node.parent_node.subtree_weight += node.subtree_weight # Find the smallest difference between the weight of a subtree and # the weight of the total tree minus that subtree. total_weight = nodes[root_node_index].subtree_weight best_diff = total_weight for node in nodes: a = node.subtree_weight b = total_weight - a diff = abs(a - b) if diff < best_diff: best_diff = diff print best_diff
to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.