[scikit-learn] How to extract subtree from a RegressionTree using the tree attribute?

[marc nicole]

I want to extract the subtree from the RegressionTree resulting from
training the associated model based on inputs: rootNode and depth,

Here's my buggy code (that I want it to be checked for errors)

def extract_tree_depth_first_traversal(tree, root_start, t_depth):
    depth = 1
    sub_tree = []
    stack = Queue()
    stack.put(root_start)
    while stack:
        current_node = stack.get(0)
        sub_tree.append(current_node)
        left_child = tree.children_left[current_node]
        if left_child >= 0:
            stack.put(left_child)
        right_child = tree.children_right[current_node]
        if right_child >= 0:
            stack.put(right_child)
        children_current_node = [left_child, right_child]
        for child in children_current_node:
            sub_tree.append(child)
        if depth >= t_depth:
            break
        depth = depth + 1
    return sub_tree

Could somebody spot the error for me ?

_______________________________________________
scikit-learn mailing list
scikit-learn@python.org
https://mail.python.org/mailman/listinfo/scikit-learn