Node Membership

Constructs a binary variable that encodes whether each observation falls into a particular node of an rpart object.

node_membership(tree, X, node)

Arguments

tree

An rpart object.

X

Covariate matrix (no intercept).

node

Number of node.

Value

Logical vector denoting whether each observation in X falls into node.

See also

subtree leaf_membership get_leaves

Author

Riccardo Di Francesco

Examples

## Generate data.
set.seed(1986)

n <- 3000
k <- 3

X <- matrix(rnorm(n * k), ncol = k)
colnames(X) <- paste0("x", seq_len(k))

Y <- exp(X[, 1]) + 2 * X[, 2] * X[, 2] > 0 + rnorm(n)

## Construct tree.
library(rpart)
tree <- rpart(Y ~ ., data = data.frame(Y, X))

## Extract number of leaves.
is_in_third_node <- node_membership(tree, X, 3)
head(is_in_third_node)
#> [1]  TRUE  TRUE  TRUE FALSE  TRUE FALSE