from collections import Counter
x = [1, 2, 3, 4, 5, 6]
y = [0, 0, 1, 1, 1, 0]
def gini(labels):
    if not labels:
        raise ValueError('empty node')
    n = len(labels)
    return 1-sum((c/n)**2 for c in Counter(labels).values())
root = gini(y)
results = []
for lo, hi in zip(x, x[1:]):
    t = (lo+hi)/2
    left = [v for u, v in zip(x, y) if u <= t]
    right = [v for u, v in zip(x, y) if u > t]
    weighted = (len(left)*gini(left)+len(right)*gini(right))/len(y)
    results.append((t, weighted))
print('root', root)
print([(t, round(g, 6)) for t, g in results])
best = min(results, key=lambda item: item[1])
print('best', best, 'gain', root-best[1])
assert best == (2.5, 0.25) and root == 0.5