ML coding · Classic algorithms
Best decision-tree split by Gini impurity
Given features X (n × d) and class labels y, find the split
"xf ≤ t goes left" with the lowest weighted Gini impurity
(nLGL + nRGR)/n, where G = 1 − Σc pc².
- Candidate thresholds for a feature are the midpoints between consecutive distinct values of that feature.
- Treat impurities within 10⁻¹² as equal. On a tie, prefer the lower feature index, then the lower threshold.
- Return
[feature, threshold, impurity]. If no feature has two distinct values, return[-1, 0.0, G]with G the impurity of the whole set.
Examples
Input X = [[1], [2], [3], [4]], y = [0, 0, 1, 1]
Output [0, 2.5, 0]
A perfect split at 2.5: both children are pure.
Input X = [[1, 5], [2, 5], [3, 6], [4, 6]], y = [0, 1, 0, 1]
Output [0, 1.5, 0.3333]
Nothing separates the classes. Thresholds 1.5 and 3.5 on feature 0 both give 1/3; the lower one wins.
Input X = [[7], [7], [7]], y = [0, 1, 1]
Output [-1, 0, 0.4444]
No threshold exists. Return −1 and the impurity of the whole set, 4/9.
Constraints
- 1 ≤ n ≤ 10⁴, 1 ≤ d ≤ 50
- Labels are integers
- Threshold and impurity checked to 10⁻⁹
Hints
Hint 1
For one feature, sort the points by that feature. A threshold between positions i and i + 1 puts the first i + 1 points on the left.
Hint 2
Keep class counts for the left side and subtract from the totals for the right. Moving one point across is an O(1) update.
Hint 3
Skip positions where the next value is equal to the current one: no threshold separates equal values.
Where candidates go wrong
- Using the data values themselves as thresholds, so that "≤ t" puts a shared value on one side or the other depending on how you wrote it.
- Recomputing the class counts from scratch for every threshold: O(n²) per feature.
- Weighting the children's impurities equally instead of by size.
- Comparing impurities with a strict < and no tolerance, so that rounding decides the tie rule.
The idea
For each feature, sort once. Sweep left to right, moving one point at a time from the right child's counts to the left child's. After moving point i, if its value differs from the next one, the midpoint between them is a valid threshold and both children's Gini values follow from the counts in O(C). The sweep is O(nC) after an O(n log n) sort, so the whole search is O(d·n(log n + C)).
Iterating features in order and thresholds in increasing order, and replacing the best only on a strict improvement beyond the tolerance, implements the tie rule without any extra code.
Time O(d · n (log n + C)) · Space O(n + C)
Reference solution
import numpy as np
def best_split(X, y):
X = np.asarray(X, dtype=float)
n, d = X.shape
classes, yi = np.unique(np.asarray(y), return_inverse=True)
total = np.bincount(yi, minlength=len(classes)).astype(float)
def gini(counts, m):
return 1.0 - ((counts / m) ** 2).sum()
best = None
for f in range(d):
order = np.argsort(X[:, f], kind='stable')
xs, ys = X[order, f], yi[order]
left = np.zeros(len(classes))
for i in range(n - 1):
left[ys[i]] += 1 # point i crosses to the left
if xs[i] == xs[i + 1]:
continue # no threshold separates equal values
nl = i + 1
g = (nl * gini(left, nl) + (n - nl) * gini(total - left, n - nl)) / n
if best is None or g < best[2] - 1e-12: # strict: earlier feature and threshold win ties
best = [f, (xs[i] + xs[i + 1]) / 2, g]
if best is None:
return [-1, 0.0, gini(total, n)]
return [best[0], float(best[1]), float(best[2])]Follow-ups an interviewer asks
- Use entropy instead of Gini. Which splits change?
- Regression trees: minimise the weighted variance of the children, with running sums of y and y².
- Build the full tree recursively with a maximum depth and a minimum leaf size.
- Histogram-based splits, as in LightGBM: bin each feature into 256 buckets first. What does it cost in accuracy?
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
Gradient-boosted trees still win most tabular problems, and quant research runs on tabular data. Interviewers ask for the split function because it is the whole tree in miniature, and because the efficient sweep separates candidates who think about complexity from those who do not.
The theory behind it is in the lesson k-means, kNN, PCA and decision trees from scratch.
Keep going
k-nearest neighbours classifier
Predict labels by majority vote among the k closest training points, with explicit tie rules.
Open the problem → Medium20 mink-means clustering (Lloyd's algorithm)
Lloyd's algorithm from given starting centroids: vectorised assignment, mean update, empty clusters and convergence, written from scratch in NumPy..
Open the problem → Medium15 minPrincipal component analysis
Centre the data, take the SVD, keep the top k directions with a fixed sign, and project.
Open the problem → Medium25 minBest decision-tree split by Gini impurity
Find the feature and threshold that minimise weighted Gini impurity, with a sorted sweep in O(n log n) per feature.
Open the problem →