Skip to content

Instantly share code, notes, and snippets.

@zpconn

zpconn/tree_weight.py

Created January 14, 2014 04:00
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save zpconn/8412833 to your computer and use it in GitHub Desktop.
Save zpconn/8412833 to your computer and use it in GitHub Desktop.
Download ZIP
Given a list of triples (a,b,c), where a is the ID of a node in a tree, b is the ID of its parent, and c is the weight of the node, this computes the total weight of every subtree. By convention, the root node has both ID and weight 0.
Raw
tree_weight.py
from collections import defaultdict
def weights(triples):
tree = defaultdict(list)
weights = defaultdict(int)
cache = defaultdict(int)
for t in triples:
weights[t[0]] = t[2]
tree[t[1]].append(t[0])
for t in triples:
print str(t[0]) + ": " + str(weight(tree, weights, cache, t[0]))
def weight(tree, weights, cache, id):
if id in cache:
return cache[id]
w = weights[id]
if id in tree:
for child in tree[id]:
w += weight(tree, weights, cache, child)
cache[id] = w
return w
print weights([(10, 30, 1), (30, 0, 10), (20, 30, 2), (50, 40, 3), (40, 30, 4)])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment