Problems
Loading...
1 / 1

Decision Tree Best Split

Classic ML
Hard

A decision tree grows by repeatedly splitting the data on the feature and threshold that best separates the classes. The quality of a split is measured by the information gain: how much the Gini impurity decreases after splitting.

Given a feature matrix X and class labels y, find the single best split (feature index and threshold) that maximizes information gain using Gini impurity.

Algorithm

  1. Compute the Gini impurity of the parent node:
Gini(S)=1kpk2\text{Gini}(S) = 1 - \sum_{k} p_k^2

Where p_k is the fraction of samples belonging to class k.

  1. For each feature and each midpoint between consecutive sorted unique values, split the data into left (feature <= threshold) and right (feature > threshold)

  2. Compute the weighted Gini impurity after the split:

Ginisplit=SLSGini(SL)+SRSGini(SR)\text{Gini}_{\text{split}} = \frac{|S_L|}{|S|} \cdot \text{Gini}(S_L) + \frac{|S_R|}{|S|} \cdot \text{Gini}(S_R)
  1. The information gain is the parent Gini minus the weighted split Gini. Return the feature and threshold with the highest gain. Break ties by smallest feature index, then smallest threshold.
Loading visualization...

Examples

Input:

X = [[1, 5], [2, 5], [3, 5], [4, 5]], y = [0, 0, 1, 1]

Output:

[0, 2.5]

Splitting on feature 0 at threshold 2.5 perfectly separates class 0 (left) from class 1 (right). The Gini impurity after this split is 0.

Input:

X = [[1, 1], [2, 1], [1, 10], [2, 10]], y = [0, 0, 1, 1]

Output:

[1, 5.5]

Feature 0 cannot separate the classes (both values appear in both classes). Feature 1 at threshold 5.5 gives a perfect split.

Hint 1

Implement a gini(labels) function: count each class, compute p_k = count_k / total, return 1 - sum(p_k^2). Then for each feature and threshold, split y into left and right, compute weighted Gini, and track the best.

Hint 2

Get sorted unique values with sorted(set(X[i][f] for i in range(n))). Thresholds are midpoints: (vals[i] + vals[i+1]) / 2. Skip if either side is empty after the split.

Requirements

  • Use Gini impurity as the splitting criterion
  • Try thresholds at the midpoint between consecutive sorted unique values for each feature
  • Weight the child Gini impurities by the fraction of samples in each child
  • Return [feature_index, threshold] as a list

Constraints

  • X has at least 2 rows and 1 column
  • y has at least 2 distinct classes
  • A valid split exists
  • Return [feature_index, threshold]
  • Time limit: 300 ms
Try Similar Problems