Improve expand_tree_rec