Gini impurity — how mixed is this pile?

DEALS = ["won"] * 6 + ["lost"] * 4

n = record("n_deals", len(DEALS))
won = record("n_won", DEALS.count("won"))
lost = record("n_lost", DEALS.count("lost"))
print(f"{won} won, {lost} lost, out of {n}")
print(f"share won  = {won}/{n} = {won / n}")
print(f"share lost = {lost}/{n} = {lost / n}")

import random
random.seed(3)

TRIALS = 100_000
disagreed = sum(1 for _ in range(TRIALS)
                if random.choice(DEALS) != random.choice(DEALS))

measured = record("measured_gini", round(disagreed / TRIALS, 3))
print(f"{disagreed} of {TRIALS} pairs disagreed  ->  {measured}")

def gini(pile):
    """1 - the chance two draws agree."""
    if not pile:
        return 0.0
    shares = [pile.count(k) / len(pile) for k in set(pile)]
    return 1 - sum(s * s for s in shares)

by_hand = 1 - (6 / 10) ** 2 - (4 / 10) ** 2
print(f"by hand      {by_hand:.3f}")
print(f"by function  {gini(DEALS):.3f}")
print(f"by playing   {measured:.3f}")

record("gini_by_hand", round(by_hand, 3))
record("gap_to_simulation", round(abs(by_hand - measured), 3))

piles = [(w, gini(["won"] * w + ["lost"] * (10 - w))) for w in range(11)]
for w, g in piles:
    print(f"{w:>2} won, {10 - w:>2} lost   {g:.2f}  {'#' * round(g * 40)}")

record("gini_5_5", round(gini(["won"] * 5 + ["lost"] * 5), 3))
record("gini_9_1", round(gini(["won"] * 9 + ["lost"]), 3))
record("gini_10_0", round(gini(["won"] * 10), 3))
record("gap_5_5_to_6_4", round(gini(["won"] * 5 + ["lost"] * 5) - by_hand, 3))

def plot(ax):
    ps = [i / 200 for i in range(201)]
    ax.plot(ps, [1 - p * p - (1 - p) ** 2 for p in ps], color="#9aa0aa", lw=1.4, zorder=1)
    ax.scatter([w / 10 for w, _ in piles], [g for _, g in piles], s=34, color="#e2574c", zorder=3)
    for w, g in [(0, 0.0), (5, 0.5), (6, by_hand)]:
        label = {0: "0 won\nnot mixed at all", 5: "5 and 5\nas mixed as it gets",
                 6: "our pile\n6 won, 4 lost"}[w]
        ax.annotate(label, (w / 10, g), textcoords="offset points",
                    xytext=(0, 14 if w == 0 else -40), ha="center", fontsize=9, color="#77777f")
    ax.set_xlabel("Share of the pile that were won")
    ax.set_ylabel("Gini impurity")
    ax.set_ylim(-0.16, 0.62)

save_fig("every-pile-of-ten", plot, figsize=(7, 4.0))

for k in range(2, 7):
    even = [str(i) for i in range(k)] * 12
    print(f"{k} outcomes, evenly spread   {gini(even):.3f}   (= 1 - {k} x (1/{k})^2)")

record("max_two", round(gini(["a", "b"] * 12), 3))
record("max_three", round(gini(["a", "b", "c"] * 12), 3))

three = ["won"] * 5 + ["lost"] * 3 + ["broke even"] * 2
record("gini_three_way", round(gini(three), 3))
print(f"\n5 won, 3 lost, 2 broke even  ->  {gini(three):.3f}")

SPLIT = {
    "budget signed off": ["won"] * 4,
    "no budget":         ["won"] * 2 + ["lost"] * 4,
}

after = 0.0
for side, pile in SPLIT.items():
    w = len(pile) / n
    print(f"{side:<20} {len(pile):>2} deals   impurity {gini(pile):.2f}   weight {w:.1f}")
    after += w * gini(pile)

record("after_split", round(after, 3))
record("plain_average", round(sum(gini(p) for p in SPLIT.values()) / len(SPLIT), 3))
print(f"\nweighted    {after:.3f}   <- what a tree compares against")
print(f"plain mean  {sum(gini(p) for p in SPLIT.values()) / len(SPLIT):.3f}   <- wrong: it lets 4 spotless deals count as much as 6 messy ones")

def plot2(ax):
    labels = ["before the split\n10 deals", "budget signed off\n5 deals", "no budget\n5 deals"]
    vals = [gini(DEALS), gini(SPLIT["budget signed off"]), gini(SPLIT["no budget"])]
    ax.bar(labels, vals, color=["#9aa0aa", "#e2574c", "#e2574c"], width=0.55)
    for i, v in enumerate(vals):
        ax.text(i, v + 0.015, f"{v:.2f}", ha="center", fontsize=10)
    ax.axhline(after, ls="--", lw=1.2, color="#3b6fd4")
    ax.text(2.42, after + 0.015, f"weighted\n{after:.2f}", ha="center", fontsize=9, color="#3b6fd4")
    ax.set_ylabel("Gini impurity")
    ax.set_ylim(0, 0.58)
    ax.grid(axis="x", visible=False)

save_fig("before-and-after-a-split", plot2, figsize=(7, 3.9))