Convert a complete binary search tree in to van emde boas order

gistfile1.py

#!/usr/bin/python
from sys import argv
from random import Random

def build_search_tree(a):
	assert(len(a))
	if len(a) == 1:
		return (a[0], None, None)
	n = len(a)
	return (a[n/2], build_search_tree(a[:n/2]),
			build_search_tree(a[(n/2)+1:]))

def layout_veb(node, h):
	if node is None:
		return None

	(me, left, right) = node
	if h == 1:
		assert(left == None)
		assert(right == None)
		return [me]

	def split(node, h, bottoms, depth = 0):
		if node is None:
			return None
		if depth >= h:
			bottoms.append(node)
			return None
		(me, left, right) = node
		return ((me, split(left, h, bottoms, depth + 1),
				split(right, h, bottoms, depth + 1)))

	h0 = h / 2
	bottoms = []
	top = split((me, left, right), h0, bottoms)
	ret = layout_veb(top, h0)
	map(lambda x:ret.extend(layout_veb(x, h - h0)), bottoms)
	return ret

def main(argv):
	EXIT_SUCCESS = 0
	EXIT_FAILURE = 1

	try:
		h = int(argv[1])
	except:
		print 'Usage:\n\t%s <num>'%argv[0]
		return EXIT_FAILURE

	num = (1 << h) - 1
	print 'height %u vEB search tree for %u items'%(h, num)
	r = Random(0)

	#items = map(lambda x:r.getrandbits(h + 1), xrange(num))
	#items.sort()
	items = range(num)
	print 'in order: ', items
	tree = build_search_tree(items)

	veb = layout_veb(tree, h)
	print 'vEB order:', veb

	return EXIT_SUCCESS

if __name__ == '__main__':
	raise SystemExit, main(argv)