python "compressed" radix tree (in plain arrays)

radix.py

def radix_tree_push_last_subtree_down(tree, level, segment_index):
	prev_word_old_segment, prev_word_old_next_child_pos = tree[level][-1]
	assert segment_index < len(prev_word_old_segment)
	prev_word_new_segment_prefix = prev_word_old_segment[:segment_index]
	prev_word_new_segment_suffix = prev_word_old_segment[segment_index:]
	tree[level][-1] = (prev_word_new_segment_prefix, prev_word_old_next_child_pos)
	old_next_level = tree[level + 1][prev_word_old_next_child_pos:]
	tree[level + 1] = tree[level + 1][:prev_word_old_next_child_pos]

	if len(tree) == level + 2:
		tree.append([])
	old_next_next_level_pos = len(tree[level + 2])

	for node in old_next_level:
		if type(node) == tuple:
			old_next_next_level_pos = node[1]
			break

	tree[level + 1].append((prev_word_new_segment_suffix, old_next_next_level_pos))
	level += 2

	new_level = old_next_level
	while True:
		if all(lambda node: type(node) == int, new_level):
			tree[level].extend(new_level)
			break
		if len(tree) == level + 1:
			tree.append([])

		old_next_next_level_pos = len(tree[level + 1])
		for node in new_level:
			if type(node) == tuple:
				old_next_next_level_pos = node[1]
				break

		old_next_level = tree[level][old_next_next_level_pos:]
		tree[level] = tree[level][:old_next_next_level_pos]

		for node in new_level:
			if type(node) == tuple:
				old_next_level_pos = node[1]
				break
		diff = len(tree[level + 1]) - old_next_level_pos
		for i, node in enumerate(new_level):
			if type(node) == tuple:
				new_level[i] = node[0], node[1] + diff
		tree[level].extend(new_level)

		new_level = old_next_level
		level += 1

	return tree

assert radix_tree_push_last_subtree_down([[('romane',0)], [1]], 0, 5) == [[('roman',0)], [('e', 0)], [1]]
assert (radix_tree_push_last_subtree_down([[('roman',0)], [('e', 0), ('us', 1)], [1, 2]], 0, 3) ==
	[
		[('rom',0)],
		[('an',0)],
		[('e', 0), ('us', 1)],
		[1, 2]
	]
)

# assumes words are inserted in dictionary order
def radix_tree_insert_word(tree, previous_word, word, value):
	level = 0
	char_index = 0
	segment_index = 0
	while char_index < len(previous_word) and previous_word[char_index] == word[char_index]:
		char_index += 1
		segment_index += 1
		if len(tree[level][-1][0]) == segment_index:
			level += 1
			segment_index = 0

	if segment_index > 0:
		radix_tree_push_last_subtree_down(tree, level, segment_index)
		level += 1

	if len(tree) == level + 1:
		tree.append([])
	tree[level].append((word[char_index:], len(tree[level + 1])))
	level += 1
	tree[level].append(value)
	return tree

assert (radix_tree_insert_word([[('romane',0)], [1]], 'romane', 'romanus', 2) ==
	[[('roman',0)], [('e', 0), ('us', 1)], [1, 2]])
assert (radix_tree_insert_word([[('roman',0)], [('e', 0), ('us', 1)], [1, 2]], 'romanus', 'romulus', 3) ==
	[
		[('rom',0)],
		[('an',0), ('ulus',2)],
		[('e',0), ('us',1), 3],
		[1, 2]
	]
)