mst for hierarchical trees

mst.py

import collections
def counts(ls, d):
  cnts = collections.defaultdict(int)
  for l in ls:
    while l:
      cnts[l] += 1
      l = d[l]
  return cnts

def mst(ls, d):
  cnts = counts(ls, d)
  vs = set()
  ss = set()
  for l in ls:
    cur_l = l_cand = l
    while cur_l:
      if cnts[cur_l] > cnts[l_cand]:
        l_cand = cur_l
      cur_l = d[cur_l]
      if cur_l in ss:
        break
      elif cur_l:
        ss.add(cur_l)
    else:
      vs.add(l_cand) 
  return vs


def visual(d):
  cs = collections.defaultdict(list)  
  for k, v in d.iteritems():
    cs[v].append(k)
  def p(i):
    if i in cs:
      return {i: map(p, cs[i])}
    return i
  return p(0)

import pprint
import random
d = dict((r, random.randint(0, r - 1)) for r in xrange(1, 200))