Create a gist now

Instantly share code, notes, and snippets.

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[0] - 1]
node_two = nodes[edge_indexes[1] - 1]
# Arrange the nodes in breadth-first order. Here we designate nodes[0]
# 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:
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment