Best Decision Tree Split

Problem

Find the (feature_index, threshold) split that minimizes the weighted Gini impurity of the two children.

For each feature, consider thresholds midway between consecutive unique values. Return [feature_index, threshold] of the best split.

If no split improves (only one unique value per feature), return [0, 0.0].

Hint

Gini of a node = 1 - sum(p_i^2). Weighted by node size.

case 1

in: [[[1,2],[2,3],[3,4],[8,9],[9,10]],[0,0,0,1,1]]

expected: [0,5.5]

case 2

in: [[[1],[2],[10],[11]],[0,0,1,1]]

expected: [0,6]

case 3

in: [[[1,1],[1,2],[2,1],[2,2]],[0,1,1,0]]

expected: [0,1.5]

Starter codePython

def best_split_gini(X, y):
    # Returns [feature_index, threshold]
    # Your code here
    pass

3 test cases · run your code against them in the editor