Code to create and query implicit van emde boas layout search trees

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)))

	def fls(i):
		ret = 0
		while i:
			i >>= 1
			ret += 1
		return ret
	def hyperceil(i):
		return 1 << fls(i - 1)

	h0 = h / 2
	h0 = hyperceil((h + 1) / 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 bfs_to_veb(bfs, h):
	def fls(i):
		ret = 0
		while i:
			i >>= 1
			ret += 1
		return ret
	def ilog2(i):
		return fls(i) - 1
	def hyperceil(i):
		return 1 << fls(i - 1)

	if h <= 2:
		return bfs

	depth = ilog2(bfs)
	split = h / 2 #hyperceil((h + 1) / 2)
	bottom_height = split
	top_height = h - bottom_height

	if depth < top_height:
		return bfs_to_veb(bfs, top_height)

	subtree_depth = depth - top_height
	subtree_root = bfs >> subtree_depth;

	num_subtrees = 1 << top_height
	bfs &= (1 << subtree_depth) - 1
	bfs |= (1 << subtree_depth)

	subtree_size = (1 << bottom_height) - 1
	toptree_size = (1 << top_height) - 1

	prior_length = toptree_size + \
			(subtree_root & (num_subtrees - 1)) * subtree_size
	return prior_length + bfs_to_veb(bfs, bottom_height)

def veb_lookup(veb, key):
	def height(n):
		ret = 0
		while n & 1:
			ret += 1
			n >>= 1
		return ret
	h = height(len(veb))
	print 'query veb: %u nodes, height %u'%(len(veb), h)
	
	n = 1
	while True:
		if n > len(veb) + 1:
			print 'done'
			return None
		idx = bfs_to_veb(n, h) - 1
		if idx >= len(veb):
			print 'done'
			return None
		print '- check key %d @ idx %d (bfs %d)'%(veb[idx], idx, n)
		if key < veb[idx]:
			n = (n << 1) + 0
		elif key > veb[idx]:
			n = (n << 1) + 1
		else:
			print 'key found @ %d'%idx
			return key;
	#for x in xrange(len(veb)):
	#	print '%d -> %d'%(x, bfs_to_veb(x, h))

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(1, num + 1)
	#print 'in order: ', items
	tree = build_search_tree(items)

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

	veb_lookup(veb, 13)
	return EXIT_SUCCESS

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