Building a Decision Tree From Scratch — Understanding the Internal Working of it

This is the third entry in my “ML from scratch” series, after KNN and Gaussian Naive Bayes. This time: a decision tree classifier, built with nothing but NumPy, broken down from the concept all the way to individual lines of code. What Is a Decision Tree? A decision tree is a model that makes predictions by asking a sequence of yes/no questions about the input’s features. Each question narrows down the possibilities until you arrive at an answer. Structurally, it’s a binary tree: Every internal node holds a question of the form “is feature X ≤ some threshold?” Every leaf node holds a predicted class label. To classify a new sample, you start at the root, answer the question at each node, follow the corresponding branch (left if true, right if false), and repeat until you land on a leaf. Whatever label that leaf holds is the prediction. That’s it structurally. The interesting part — and the part that actually needs an algorithm — is how the tree decides which questions to ask, and in what order. Simple flow of how decision tree executes How It Works: The Core Idea Training a decision tree means building this tree of questions from data, one node at a time, top-down. At the root, you have your entire training set, with a mix of classes. The goal is to find one question — one (feature, threshold) pair — that splits the data into two groups that are, as much as possible, more homogeneous than the group you started with. Ideally, one side ends up mostly class A, the other mostly class B. Once you’ve picked that first question and split the data, you don’t stop — you treat each of the two resulting groups as its own smaller problem, and repeat the same process on each: find the best question to split that group further. This continues recursively, each split producing two more (smaller) groups to potentially split again. The recursion stops when further splitting stops making sense — a group is already pure (all one class), you’ve split enough times already, or there aren’t enough samples left to justify splitting further. At that point, the group becomes a leaf, and the leaf’s prediction is simply the majority class within it. So the whole algorithm is really just two ideas layered together: A way to score how good a candidate split is (so you can pick the best one at each node). Recursion — apply that scoring process over and over, on smaller and smaller subsets, until you’re out of useful splits to make. Everything else is implementation detail. The next section covers idea #1 — how “good” gets defined mathematically. Methodology: Entropy and Information Gain To score a split, we need a way to measure how mixed up (impure) a set of labels is — and how much a candidate split reduces that impurity. Decision trees commonly use entropy for the first part and information gain for the second. Entropy Entropy measures the disorder in a set of labels: E(S) = -Σ p(x) · log(p(x)) Where p(x) is the proportion of class x in set S, summed over all classes present. Worked example. A node with 10 samples: 6 of class 0, 4 of class 1. p(0) = 6/10 = 0.6 p(1) = 4/10 = 0.4 E = -(0.6 · log(0.6) + 0.4 · log(0.4)) = -(0.6 · (-0.511) + 0.4 · (-0.916)) = -(-0.3065 + -0.3665) = 0.673 Compare that to a pure node — 10 samples, all class 0: p(0) = 1.0 E = -(1.0 · log(1.0)) = -(1.0 · 0) = 0 Entropy of 0 means no disorder at all — nothing left to gain from splitting further. Entropy is at its maximum when classes are perfectly balanced, and shrinks toward 0 as one class comes to dominate the set. In code: def _entropy(self, y): hist = np.bincount(y) ps = hist / len(y) return -np.sum([p * np.log(p) for p in ps if p > 0]) Information Gain Entropy alone tells you how mixed a single set is. To evaluate a split , you compare the parent’s entropy against a weighted average of the two children’s entropy: IG = E(parent) - [ (n_left/n) · E(left) + (n_right/n) · E(right) ] Weighting by n_left/n and n_right/n matters: a split that carves off a large, pure group and leaves a small, mixed remainder should score differently than one that produces two medium, moderately-mixed groups. Weighting by size accounts for that. Worked example , continuing the 10-sample node (parent entropy E = 0.673). Suppose a candidate split sends 5 samples left (all class 0) and 5 right (1 class-0, 4 class-1): E(left) = 0 (pure) E(right) = -(0.2·log(0.2) + 0.8·log(0.8)) = -(0.2·(-1.609) + 0.8·(-0.223)) = -(-0.322 + -0.179) = 0.500 Weighted child entropy = (5/10)·0 + (5/10)·0.500 = 0.250 IG = 0.673 - 0.250 = 0.423 That’s a strong split — it fully isolated a pure group. A weak split, where the class ratio barely changes on either side, produces an IG close to 0. In code: def _information_gain(self, y, X_column, threshold): parent_entropy = self._entropy(y) left_idxs, right_idxs = self._split(X_column, threshold) if len(left_idxs) == 0 or len(right_idxs) == 0: return 0 n = len(y) n_l, n_r = len(left_idxs), len(right_idxs) e_l, e_r = self._entropy(y[left_idxs]), self._entropy(y[right_idxs]) child_entropy = (n_l/n) * e_l + (n_r/n) * e_r return parent_entropy - child_entropy Choosing the Best Split With a way to score any single candidate split, finding the best split at a node is a search: try every feature, and for each feature, try every value that appears in that feature’s column as a candidate threshold, keeping whichever (feature, threshold) pair produced the highest information gain. best_split(X, y) = argmax over all (feature, threshold) of IG(y, X[:, feature], threshold) Only unique observed values need to be tried as thresholds — any value strictly between two consecutive observed values produces an identical split, so there’s no benefit to checking a finer-grained range. Walking Through the Code, Part by Part Node class Node: def __init__(self, feature=None, threshold=None, left=None, right=None, *, value=None): self.feature = feature self.threshold = threshold self.right = right self.left = left self.value = value def is_leaf_node(self): return self.value is not None Node is deliberately dual-purpose — the same class represents both internal nodes and leaves, distinguished only by which fields are set: An internal node has feature and threshold (the question it asks) plus left and right (its two children). value stays None. A leaf node has only value set (the predicted class). Everything else stays None. value is keyword-only (the * forces this) specifically so it can't be passed positionally by accident and confused with left/right — leaves and internal nodes are constructed with visually distinct calls: Node(value=leaf_value) vs. Node(best_feature, best_threshold, left, right). is_leaf_node() checks self.value is not None — which works because internal nodes never set value, and leaves never set anything else. This one check is what predict uses to decide whether to stop traversing or keep going. __init__ def __init__(self, min_samples_split=2, max_depth=100, n_features=None): self.min_samples_split = min_samples_split self.max_depth = max_depth self.n_features = n_features self.root = None Three hyperparameters, all controlling when the tree stops growing (directly or indirectly), plus self.root, which starts empty and gets filled in by fit. min_samples_split — the minimum number of samples a node needs to be eligible for splitting at all. max_depth — a hard cap on how many splits deep the tree can go. n_features — how many features to consider at each split. None means "use all of them"; a smaller number introduces the random feature subsampling used by Random Forests. fit def fit(self, X, y): self.n_features = X.shape[1] if not self.n_features else min(X.shape[1], self.n_features) self.root = self.grow_tree(X, y) The entry point. The first line resolves n_features into an actual number: if none was specified, use every feature in X; if one was specified, use whichever is smaller — the requested count or the number of features actually available (so you can't accidentally ask for more features than exist). The second line kicks off recursive tree-building and stores the resulting root node. grow_tree def grow_tree(self, X, y, depth=0): n_samples, n_feats = X.shape n_labels = len(np.unique(y)) if depth >= self.max_depth or n_labels == 1 or n_samples < self.min_samples_split: leaf_value = self.most_common_label(y) return Node(value=leaf_value) feat_idxs = np.random.choice(n_feats, self.n_features, replace=False) best_feature, best_threshold = self.best_split(X, y, feat_idxs) left_idxs, right_idxs = self._split(X[:, best_feature], best_threshold) left = self.grow_tree(X[left_idxs, :], y[left_idxs], depth + 1) right = self.grow_tree(X[right_idxs, :], y[right_idxs], depth + 1) return Node(best_feature, best_threshold, left, right) Each call handles one node, given the slice of data (X, y) that reached it and how deep it is (depth). Stopping check first. The if condition covers the three ways a node becomes a leaf: it's too deep (depth >= max_depth), it's already pure (n_labels == 1), or it has too few samples to keep splitting (n_samples < min_samples_split). If any is true, skip straight to most_common_label(y) and return a leaf — no point searching for a split that won't be used. Otherwise, split and recurse. np.random.choice(n_feats, self.n_features, replace=False) picks which features this node is even allowed to consider — all of them by default, a random subset if n_features was restricted. best_split searches those features for the best (feature, threshold) pair. _split then partitions the actual data into left_idxs/right_idxs based on that choice. The two recursive calls — self.grow_tree(X[left_idxs, :], y[left_idxs], depth + 1) and the equivalent for right — are where the "repeat the whole process on each smaller group" idea from earlier actually happens in code. Each call returns a fully built subtree (root node of that subtree), and the final line wires both of those subtrees into a new internal Node, which is what gets returned up to this call's caller. That's how depth-first construction bubbles all the way back up to a single root. best_split def best_split(self, X, y, feat_idxs): best_gain = -1 split_idx, split_threshold = None, None for feat_idx in feat_idxs: X_column = X[:, feat_idx] thresholds = np.unique(X_column) for thr in thresholds: gain = self._information_gain(y, X_column, thr) if gain > best_gain: best_gain = gain split_idx = feat_idx split_threshold = thr return split_idx, split_threshold A brute-force search, directly implementing the argmax from the methodology section. The outer loop goes feature by feature; the inner loop goes threshold by threshold (every unique value observed in that feature's column). For every (feature, threshold) pair, _information_gain scores it, and best_gain/split_idx/split_threshold track the best one seen so far. best_gain starts at -1 specifically because information gain is always ≥ 0 — so the very first real split evaluated is guaranteed to beat the initial placeholder and get recorded. _information_gain and _split def _information_gain(self, y, X_column, threshold): parent_entropy = self._entropy(y) left_idxs, right_idxs = self._split(X_column, threshold) def _split(self, X_column, split_thresh): left_idxs = np.argwhere(X_column <= split_thresh).flatten() right_idxs = np.argwhere(X_column > split_thresh).flatten() return left_idxs, right_idxs These two are the direct code form of the entropy/information-gain formulas covered above — _split produces the two index arrays for "at or below the threshold" vs. "above it" using np.argwhere and a boolean comparison; _information_gain uses those indices to compute each side's entropy and combine them into the weighted score. The one piece of defensive logic — if len(left_idxs) == 0 or len(right_idxs) == 0: return 0 — handles a threshold that doesn't actually separate anything (every sample lands on one side), which would otherwise divide by an empty array when computing that side's entropy. _entropy and most_common_label def _entropy(self, y): hist = np.bincount(y) ps = hist / len(y) return -np.sum([p * np.log(p) for p in ps if p > 0]) def most_common_label(self, y): counter = Counter(y) return counter.most_common(1)[0][0] _entropy is the formula from the methodology section, translated line for line: np.bincount(y) counts samples per class, dividing by len(y) turns those into proportions, and the list comprehension sums p · log(p) over every class with p > 0 (skipping zero-count classes, since log(0) is undefined and they contribute nothing anyway). most_common_label is what a leaf actually predicts: Counter(y).most_common(1) returns the single most frequent label in y as a (label, count) tuple, and [0][0] pulls out just the label. predict and _traverse_tree def predict(self, X): return np.array([self._traverse_tree(x, self.root) for x in X]) def _traverse_tree(self, x, node): if node.is_leaf_node(): return node.value if x[node.feature] <= node.threshold: return self._traverse_tree(x, node.left) return self._traverse_tree(x, node.right) predict just runs _traverse_tree on every row of X and collects the results into an array. _traverse_tree is the inference-time mirror of the "start at the root, answer the question, follow the branch" description from the intro: check if the current node is a leaf (if so, return its stored label — done); otherwise, compare x[node.feature] against node.threshold and recurse into node.left or node.right accordingly. No computation happens here — training already decided every question in advance, so prediction is pure traversal, taking O(depth) steps per sample regardless of how large the training set was. Results Trained and evaluated on sklearn.datasets.load_breast_cancer (80/20 split, random_state=1234, default hyperparameters): Accuracy: 94% That’s competitive with sklearn’s own DecisionTreeClassifier on the same split — expected, since the underlying math is identical. sklearn's version is faster thanks to more optimized split-finding, and supports additional criteria (Gini impurity) and pruning options this implementation doesn't. To go beyond the raw accuracy number, I added a few visualizations: a confusion matrix to see where the errors landed, a feature-importance chart based on how often each feature was used to split (which features the tree actually relied on), and a PCA projection of the test set colored by correct vs. incorrect predictions, to check whether misclassifications clustered in any particular region of the data. What This Implementation Doesn’t Do No pruning. A fully-grown tree (especially with max_depth=100) can overfit badly on noisier data. Real implementations add pre-pruning (stricter stopping criteria) or post-pruning (cost-complexity pruning, ccp_alpha in sklearn). No Gini impurity option. Entropy is one valid splitting criterion; Gini impurity is another, cheaper to compute (no logarithms) and commonly used as a default in practice. Exhaustive threshold search. Trying every unique value in every feature works fine for a learning implementation, but scales poorly — sklearn’s C-optimized version does the same search far more efficiently. No built-in feature importance. The version I computed for the plots is a simple split-frequency count, not the sample-weighted (samples × information gain) importance sklearn reports. Next Steps The natural extension from here is a Random Forest — since n_features subsampling is already built in, most of the ensemble machinery is really just: train many of these trees on bootstrapped samples of the data, and aggregate their predictions by majority vote. That's next in this series. Complete Code of Decision Tree implementation is available on my GitHub. GitHub link — https://github.com/Archan47/Machine-Learning-Algorithms-From-Scratch Code for this and the rest of the “ML from scratch” series is on GitHub. Building a Decision Tree From Scratch — Understanding the Internal Working of it was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.