Constructs a variable that encodes in which leaf of an rpart object the units in a given data frame fall.

leaf_membership(tree, X)

Arguments

tree

An rpart object.

X

Covariate matrix (no intercept).

Value

A factor whose levels denote in which leaf each unit falls. Leaves are ordered in increasing order of their predictions (from most negative to most positive).

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.
leaves_factor <- leaf_membership(tree, X)
head(leaves_factor)
#> 1 2 3 4 5 6 
#> 5 7 7 2 7 2 
#> Levels: 1 2 3 4 5 6 7 8