From Scratch: K-Means and Trees
No libraries, no magic.
A decision tree feels like a smart, complicated thing. But it's really just two small ideas, repeated. Find the best question, then do the same again on each side. That's recursion, and you're about to write it. First, a quick warm-up with k means.
- Trees feel complicated
- Two small ideas, repeated
- First, a warm-up
K-means from scratch
K means is short. We pick K random points as the first centres. Then we loop. This line computes the distance from each point to each centre, using NumPy's broadcasting. Argmin assigns each point to its nearest centre. Next, each centre moves to the mean of its group. If nothing has moved, we stop. Run it. It stops after a handful of rounds, with three clean groups. That's the assign and update loop that we saw animated, in only ten lines.
import numpy as np
from sklearn.datasets import make_blobs
def kmeans(X, k, seed=0):
rng = np.random.default_rng(seed)
centres = X[rng.choice(len(X), k, replace=False)]
for round_ in range(100):
dist = ((X[:, None] - centres) ** 2).sum(axis=2)
labels = dist.argmin(axis=1)
new = np.array([X[labels == j].mean(axis=0) for j in range(k)])
if np.allclose(new, centres):
break
centres = new
return labels, centres, round_
X, _ = make_blobs(n_samples=500, centers=3, random_state=4)
labels, centres, rounds = kmeans(X, 3)
print("Stopped after", rounds, "rounds")
print("Group sizes:", np.bincount(labels))
print("Centres:", centres.round(1).tolist())Stopped after 6 rounds Group sizes: [160 166 174] Centres: [[9.5, 0.9], [4.0, -5.7], [9.5, 4.3]]
A decision tree from scratch
Now the tree. Gini measures impurity. Best split tries several thresholds on each feature, and scores each split by the weighted impurity of the two sides. Grow is the recursive heart. Find the best split, and if we've reached the depth limit, return a leaf with the majority class. If not, split the data, and call grow on each side. Walk follows the questions down to a leaf. Run it. A real decision tree, with no library, and over ninety percent accuracy.
import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
def gini(y):
p = np.bincount(y) / len(y)
return 1 - (p ** 2).sum()
def best_split(X, y):
best = (None, None, gini(y))
for f in range(X.shape[1]):
for t in np.percentile(X[:, f], [20, 40, 60, 80]):
left = X[:, f] <= t
if left.all() or not left.any():
continue
score = (left.mean() * gini(y[left])
+ (1 - left.mean()) * gini(y[~left]))
if score < best[2]:
best = (f, t, score)
return best
def grow(X, y, depth):
f, t, _ = best_split(X, y)
if depth == 0 or f is None:
return int(np.bincount(y).argmax())
left = X[:, f] <= t
return (f, t, grow(X[left], y[left], depth - 1),
grow(X[~left], y[~left], depth - 1))
def walk(node, row):
while not isinstance(node, int):
f, t, left, right = node
node = left if row[f] <= t else right
return node
X, y = load_breast_cancer(return_X_y=True)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, random_state=0)
tree = grow(X_tr, y_tr, depth=4)
pred = np.array([walk(tree, r) for r in X_te])
print("Accuracy:", (pred == y_te).mean().round(3))Accuracy: 0.93
A forest from scratch
Once you can grow a tree, a forest is easy. For each of one hundred trees, sample rows with replacement. Then fit a tree, which considers a random subset of features at each split. To predict, average the votes. Run it. The forest beats one tree. Bagging, random features, and a vote. That's the entire idea, in about ten lines.
import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
X, y = load_breast_cancer(return_X_y=True)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, random_state=0)
rng = np.random.default_rng(0)
trees = []
for i in range(100):
rows = rng.integers(0, len(X_tr), len(X_tr))
t = DecisionTreeClassifier(max_features="sqrt", random_state=i)
trees.append(t.fit(X_tr[rows], y_tr[rows]))
votes = np.mean([t.predict(X_te) for t in trees], axis=0)
single = (trees[0].predict(X_te) == y_te).mean()
forest = ((votes >= 0.5) == y_te).mean()
print("One tree:", single.round(3))
print("Forest: ", forest.round(3))One tree: 0.909 Forest: 0.965
Summary
You've now built K N N, logistic regression, k means, a decision tree, and a random forest, as well as gradient descent and boosting. There's no magic left. Next, we'll look at some practical situations that textbooks ignore, starting with what to do when you have very few labels.
Key takeaways- K-means: assign, update
- Tree: best split, recurse
- Forest: sample, grow, vote