df <- read.csv("https://roualdes.sfo3.digitaloceanspaces.com/data/abalone.csv") mf <- model.frame(age ~ 0 + length + diameter + height + whole_weight + shucked_weight + viscera_weight + shell_weight, data = df) X <- model.matrix(mf, data = df) y <- model.response(mf) find_best_split <- function(X_, y_, idx) { Xx <- X_[idx,] yx <- y_[idx] J <- ncol(Xx) N <- nrow(Xx) S <- apply(Xx, 2, \(x) quantile(x, 1:19/20)) Srows <- nrow(S) Scols <- ncol(S) predictions <- matrix(NA, nrow = Srows, ncol = Scols) for (j in 1:J) { xj <- Xx[,j] for (i in 1:Srows) { s <- S[i, j] sdx <- xj <= s predictions[i, j] <- mean(yx[sdx], na.rm=TRUE) + mean(yx[!sdx], na.rm=TRUE) } } cdx <- which.min(predictions) mdx <- arrayInd(cdx, dim(predictions)) best <- NULL if (length(mdx) > 0) { best$split <- S[mdx] best$var <- mdx[,2] } best } left_id <- function(i) { 2L * i } right_id <- function(i) { 2L * i + 1L } parent_id <- function(i) { i %/% 2L } create_node <- function(id, is_leaf, j, s, γ, n, depth) { data.frame( id = id, is_leaf = is_leaf, split_var = j, split_value = s, prediction = γ, n = n, depth = depth, stringsAsFactors = FALSE ) } should_stop <- function(idx, depth, y) { (depth >= 2) || (length(idx) < 5) } build_tree <- function(X, y, idx, nodes, id = 1L, depth = 0L) { L <- length(nodes) if (should_stop(idx, depth, y)) { nodes[[L + 1L]] <- create_node(id, TRUE, NA_real_, NA_real_, mean(y[idx]), L, depth) return(nodes) } best <- find_best_split(X, y, idx) if (is.null(best)) { nodes[[L + 1L]] <- create_node(id, TRUE, NA_real_, NA_real_, mean(y[idx]), L, depth) return(nodes) } nodes[[L + 1L]] <- create_node(id, FALSE, best$var, best$split, NA_real_, L, depth) left_idx <- idx[X[idx, best$var] <= best$split] right_idx <- idx[X[idx, best$var] > best$split] nodes <- build_tree(X, y, left_idx, nodes = nodes, id = left_id(id), depth = depth + 1L) nodes <- build_tree(X, y, right_idx, nodes = nodes, id = right_id(id), depth = depth + 1L) nodes } predict_one <- function(tree, x) { i <- 1L repeat { tdx <- tree$id == i row <- tree[tdx, ] if (row$is_leaf) { return(row$prediction) } if (x[row$split_var] <= row$split_value) { i <- left_id(i) } else { i <- right_id(i) } } } predict <- function(tree, newdata, ...) { vapply(seq_len(nrow(newdata)), function(j) { predict_one(tree, newdata[j, ]) }, numeric(1)) } nodes <- build_tree(X, y, 1:length(y), list()) tree <- do.call(rbind, nodes) sqrt(mean((y - predict(tree, X))^2)) library(rpart) rtree <- rpart(age ~ length + diameter + height + whole_weight + shucked_weight + viscera_weight + shell_weight, data = df) sqrt(mean((y - predict(rtree))^2))