Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

scikit-learn: Iterating over nodes of DecisionTreeClassifier

 from sklearn.datasets import load_iris
 from sklearn.tree import DecisionTreeClassifier

 iris = load_iris()
 X    = iris.data
 y    = iris.target
  
 clf  = DecisionTreeClassifier()
 clf  = clf.fit(iris.data,iris.target)

How can I iterate over the nodes of clf. I couldn't find it anywhere in the documentation.

like image 905
memecs Avatar asked Dec 04 '25 19:12

memecs


1 Answers

Now, there exists an example on how this can be done in the documentation.

There, they iterate the tree using

n_nodes = clf.tree_.node_count
children_left = clf.tree_.children_left
children_right = clf.tree_.children_right
feature = clf.tree_.feature
threshold = clf.tree_.threshold

node_depth = np.zeros(shape=n_nodes, dtype=np.int64)
is_leaves = np.zeros(shape=n_nodes, dtype=bool)
stack = [(0, 0)]  # start with the root node id (0) and its depth (0)
while len(stack) > 0:
    # `pop` ensures each node is only visited once
    node_id, depth = stack.pop()
    node_depth[node_id] = depth

    # If the left and right child of a node is not the same we have a split
    # node
    is_split_node = children_left[node_id] != children_right[node_id]
    # If a split node, append left and right children and depth to `stack`
    # so we can loop through them
    if is_split_node:
        stack.append((children_left[node_id], depth + 1))
        stack.append((children_right[node_id], depth + 1))
    else:
        is_leaves[node_id] = True

print("The binary tree structure has {n} nodes and has "
      "the following tree structure:\n".format(n=n_nodes))
for i in range(n_nodes):
    if is_leaves[i]:
        print("{space}node={node} is a leaf node.".format(
            space=node_depth[i] * "\t", node=i))
    else:
        print("{space}node={node} is a split node: "
              "go to node {left} if X[:, {feature}] <= {threshold} "
              "else to node {right}.".format(
                  space=node_depth[i] * "\t",
                  node=i,
                  left=children_left[i],
                  feature=feature[i],
                  threshold=threshold[i],
                  right=children_right[i]))
like image 151
André Avatar answered Dec 07 '25 20:12

André