Implementation of binary-uplift algorithm for solving LCA problem

lca.py

# This code is in public domain.
# Written by Andrew Shulayev, 2012

from pprint import pprint

test_tree = ([0, 0, 1, 0, 2, 3, 0, 3, 2, 4, 5, 6, 8], 0)

def print_tree(tree):
  par, root = tree
  n = len(par)
  def print_node(node, tab):
    print tab, node
    for x in range(n):
      if x != root and par[x] == node:
        print_node(x, tab + '  ')
  print_node(root, '')

print_tree(test_tree)

# shitty implementation of binary-uplift algorithm
# seems like it works only when par[x] < x
# I don't even want to test it properly.

def build_table(tree):
  par, root = tree
  n = len(par)
  result = []
  for node in range(n):
    row = [par[node]]
    result.append(row)
    h = 1
    while 2**(h - 1) < n:
      row.append(result[row[-1]][h - 1])
      h += 1
  return result

test_table = build_table(test_tree)

for i, row in enumerate(test_table):
  print i, row

def build_heights(tree):
  par, root = tree
  n = len(par)
  result = [None] * n
  result[root] = 0
  while None in result:
    for node in range(n):
      if result[node] is not None or result[par[node]] is None:
        continue
      result[node] = result[par[node]] + 1
  return result

test_heights = build_heights(test_tree)

def lca(tree, table, heights, u, v):
  par, root = tree
  n = len(par)
  h = len(table[0])
  if heights[v] < heights[u]:
    v, u = u, v
  while h > 0:
    h -= 1
    if heights[v] - 2**h >= heights[u]:
      v = table[v][h]
  if u == v:
    return u
  # print u, v
  left, right = -1, len(table[0])
  while left + 1 < right:
    middle = (left + right) / 2
    # print u, v, middle
    if table[v][middle] == table[u][middle]:
      right = middle
    else:
      left = middle
  h = left
  while h >= 0:
    # print u, v, h
    if table[v][h] != table[u][h]:
      v = table[v][h]
      u = table[u][h]
    h -= 1
  return par[v]

def test(u, v, out):
  result = lca(test_tree, test_table, test_heights, u, v)
  if result == out:
    print 'ok, lca(%d, %d) = %d' % (u, v, out)
  else:
    print 'wa, lca(%d, %d) = %d, expected %d' % (u, v, result, out)

test(9, 12, 2)
test(12, 9, 2)
test(7, 2, 0)
test(10, 11, 0)
test(5, 10, 5)
test(11, 9, 0)