reingold_naive.py

from gen import Tree, gentree
from operator import lt, gt
from sys import stdout


class DrawTree:
    def __init__(self, tree, depth=-1):
        self.x = -1
        self.y = depth
        self.tree = tree
        self.children = []


# traverse to the bottom of the tree, and place the leaves at an arbitrary
#   x coordinate
# if the node is a parent, draw its subtrees, then shift the right one as close
#   to the left as possible
# place the parent in the middle of the two trees.


def reingold_tilford(tree, depth=0):
    dt = DrawTree(tree, depth)
    if len(tree) == 0:
        dt.x = 0
        return dt

    if len(tree) == 1:
        dt.children = [reingold_tilford(tree[0], depth + 1)]
        dt.x = dt.children[0].x
        return dt

    left = reingold_tilford(tree[0], depth + 1)
    right = reingold_tilford(tree[1], depth + 1)

    dt.children = [left, right]
    dt.x = fix_subtrees(left, right)
    return dt


# place the right subtree as close to the left subtree as possible


def fix_subtrees(left, right):
    wl = contour(left, lt)
    wr = contour(right, gt)
    diff = max(map(lambda x, y: x - y, wl, wr)) + 1
    # stick to the integers
    diff += (right.x + diff + left.x) % 2
    addtotree(right, diff)
    return (left.x + right.x) / 2


def addtotree(tree, val):
    tree.x += val
    if len(tree.children):
        for child in tree.children[:2]:
            addtotree(child, val)


def contour(tree, comp, level=0, cont=None):
    if not level:
        cont = [tree.x]
    elif len(cont) < level + 1:
        cont.append(tree.x)
    elif comp(cont[level], tree.x):
        cont[level] = tree.x

    for child in tree.children[:2]:
        contour(child, comp, level + 1, cont)

    return cont


# given an array of nodes, print them out reasonably on one line


def printrow(level):
    x = dict((t.x, t.tree) for t in level)
    for i in range(max(x.keys()) + 1):
        try:
            stdout.write(str(x[i])[:4])
        except:
            stdout.write("    ")


def p(tree):
    level = [tree]
    while 1:
        newlevel = []
        printrow(level)
        for t in level:
            # stdout.write("%s%s%s" % (t.x*"____", str(t)[:3], t.x))
            newlevel.extend(t.children[:2])
        print
        if not newlevel:
            break
        level = newlevel


# root = gentree("/Users/llimllib/Movies")
# root.children.reverse()
# drawtree = reingold_tilford(root)
# p(drawtree)