# Build a Decision Tree
# Core implementation example — ackleywill.com

from collections import Counter, defaultdict

# Illustrative categorical fruit examples: the lab uses multiway splits.
fruits = [
    {"label": "apple", "color": "red", "shape": "round"},
    {"label": "apple", "color": "green", "shape": "round"},
    {"label": "orange", "color": "orange", "shape": "round"},
    {"label": "orange", "color": "yellow", "shape": "round"},
    {"label": "banana", "color": "yellow", "shape": "long"},
    {"label": "banana", "color": "green", "shape": "long"},
]

def gini(rows):
    return 1 - sum((n / len(rows)) ** 2
                   for n in Counter(r["label"] for r in rows).values())

for feature in ["color", "shape"]:
    groups = defaultdict(list)
    for row in fruits:
        groups[row[feature]].append(row)
    weighted = sum(len(g) / len(fruits) * gini(g) for g in groups.values())
    print(feature, "weighted impurity:", weighted,
          "reduction:", gini(fruits) - weighted)