Decision Tree Stump
Hard
machine-learning
OpenAI
DeepMind
Implement a decision tree stump (depth-1 tree) for binary classification. Given a dataset with one numeric feature and binary labels (0 or 1), find the best threshold to split on that minimizes Gini impurity.
Return a function that classifies a value as 0 or 1, along with the threshold.
Example
features = [1, 2, 3, 4, 5]
labels = [0, 0, 1, 1, 1]
threshold, left_label, right_label = decision_stump(features, labels)
# threshold=2.5, left_label=0, right_label=1
Test Cases
Python Editor
Output
Click "Run" to see results...