TreeDiff.py

#!/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[0] - 1]
	node_two = nodes[edge_indexes[1] - 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[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:
			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