Tree depth and leaf size — knowing when to stop

QUESTIONS = {
    "star":   "Is there a star the audience turns up for?",
    "summer": "Did it open in summer?",
    "sequel": "Is it a sequel?",
    "wide":   "Did it open on 3,000 screens or more?",
    "budget": "Did it cost over $100 million?",
}
SHORT = {"star": "A star?", "summer": "Summer?", "sequel": "A sequel?",
         "wide": "3,000+ screens?", "budget": "Over $100m?"}
KEYS = list(QUESTIONS)

#                          star summer sequel wide budget  hit
FILMS = {
    "Harbour Lights":       (1,   1,     0,     0,   0,     1),
    "Ironwake II":          (1,   1,     1,     1,   0,     1),
    "The Quiet Ledger II":  (1,   0,     1,     0,   0,     1),
    "Nightfall Divide":     (1,   1,     0,     1,   1,     1),
    "Ironwake III":         (1,   0,     1,     1,   1,     1),
    "Cinder Coast II":      (1,   0,     1,     1,   1,     1),
    "Redline Returns":      (1,   0,     1,     1,   1,     1),
    "Paper Kingdoms":       (1,   1,     0,     1,   0,     1),
    "Glass Monsoon":        (1,   1,     0,     0,   1,     0),
    "Cinder Coast III":     (1,   0,     1,     1,   0,     0),
    "Vermilion Rising II":  (0,   0,     1,     1,   1,     1),
    "Saltwater Sunday":     (0,   1,     0,     1,   1,     1),
    "Field of Static II":   (0,   0,     1,     0,   0,     1),
    "Wildflower County":    (0,   1,     0,     1,   0,     1),
    "Tin Sky":              (0,   1,     0,     0,   0,     0),
    "The Cartographer":     (0,   0,     0,     0,   0,     0),
    "Neon Bazaar":          (0,   1,     0,     0,   1,     0),
    "Ash & Ivory II":       (0,   0,     1,     1,   0,     0),
    "Slow River":           (0,   0,     0,     0,   0,     0),
    "The Understudy":       (0,   1,     0,     0,   0,     0),
    "Meridian":             (0,   0,     0,     1,   1,     0),
    "Copper Harbour":       (0,   1,     0,     1,   1,     0),
    "The Winter Post":      (0,   0,     0,     0,   0,     0),
    "Little Eden":          (0,   1,     0,     0,   1,     0),
}
ALL = list(FILMS)

def answers(f): return FILMS[f][:len(KEYS)]
def hit(f): return FILMS[f][-1]

def split(films, q):
    i = KEYS.index(q)
    return ([f for f in films if answers(f)[i]], [f for f in films if not answers(f)[i]])

def gini(films):
    if not films:
        return 0.0
    p = sum(hit(f) for f in films) / len(films)
    return 1 - p * p - (1 - p) * (1 - p)

def gain(films, q):
    yes, no = split(films, q)
    if not yes or not no:
        return 0.0
    return gini(films) - (len(yes) * gini(yes) + len(no) * gini(no)) / len(films)

print(f"{len(ALL)} films, {sum(hit(f) for f in ALL)} hits, impurity {gini(ALL):.2f}")

star_films, no_star = split(ALL, "star")

print(f"{'question':<18} {'all 24':>8} {'the 10 with a star':>20} {'the 14 without':>16}")
for q in KEYS:
    print(f"{SHORT[q]:<18} {gain(ALL, q):>8.3f} {gain(star_films, q):>20.3f} {gain(no_star, q):>16.3f}")

best_star = max(KEYS, key=lambda q: (gain(star_films, q), q))
best_no_star = max(KEYS, key=lambda q: (gain(no_star, q), q))
record("n_star_films", len(star_films))
record("star_hits", sum(hit(f) for f in star_films))
record("n_no_star", len(no_star))
record("no_star_hits", sum(hit(f) for f in no_star))
record("best_on_star", QUESTIONS[best_star])
record("best_on_no_star", QUESTIONS[best_no_star])
record("sequel_on_star", round(gain(star_films, "sequel"), 3))
record("wide_on_star", round(gain(star_films, "wide"), 3))
record("sequel_on_no_star", round(gain(no_star, "sequel"), 3))
record("wide_on_no_star", round(gain(no_star, "wide"), 3))
# A question already asked is worth exactly nothing below itself. The tree is
# never told this; it falls out of the arithmetic, because one side is empty.
record("star_gain_below_itself", round(gain(star_films, "star"), 3))

def plot(ax):
    xs = range(len(KEYS))
    ax.bar([x - 0.25 for x in xs], [gain(ALL, q) for q in KEYS], width=0.24,
           color="#9aa0aa", label="all 24 films")
    ax.bar([x for x in xs], [gain(star_films, q) for q in KEYS], width=0.24,
           color="#e2574c", label="the 10 with a star")
    ax.bar([x + 0.25 for x in xs], [gain(no_star, q) for q in KEYS], width=0.24,
           color="#3b6fd4", label="the 14 without")
    ax.set_xticks(list(xs))
    ax.set_xticklabels([SHORT[q] for q in KEYS], fontsize=9)
    ax.set_ylabel("Impurity removed")
    ax.legend(frameon=False, fontsize=9)
    ax.grid(axis="x", visible=False)

save_fig("a-different-question-wins", plot, figsize=(7, 3.8))

def build(films, min_leaf=1, max_depth=None, depth=0):
    call = "hit" if sum(hit(f) for f in films) * 2 >= len(films) else "flop"
    note = f"{len(films)} films - {sum(hit(f) for f in films)} hit"
    stop = (gini(films) == 0
            or (max_depth is not None and depth >= max_depth)
            or len(films) < 2 * min_leaf)
    if not stop:
        usable = [q for q in KEYS
                  if all(len(side) >= min_leaf for side in split(films, q)) and gain(films, q) > 0]
        if usable:
            best = max(usable, key=lambda q: (gain(films, q), q))
            yes, no = split(films, best)
            return {"q": SHORT[best], "note": note,
                    "yes": build(yes, min_leaf, max_depth, depth + 1),
                    "no": build(no, min_leaf, max_depth, depth + 1)}
    return {"leaf": call, "note": note}

def leaves(node):
    return [node] if "leaf" in node else leaves(node["yes"]) + leaves(node["no"])

def depth_of(node):
    return 0 if "leaf" in node else 1 + max(depth_of(node["yes"]), depth_of(node["no"]))

def predict(node, f):
    while "leaf" not in node:
        i = KEYS.index(next(k for k in KEYS if SHORT[k] == node["q"]))
        node = node["yes"] if answers(f)[i] else node["no"]
    return node["leaf"]

def score(node):
    right = sum(1 for f in ALL if (predict(node, f) == "hit") == bool(hit(f)))
    return right, len(ALL)

FULL = build(ALL)
right, total = score(FULL)
record("full_depth", depth_of(FULL))
record("full_leaves", len(leaves(FULL)))
record("full_singletons", sum(1 for l in leaves(FULL) if l["note"].startswith("1 films")))
record("full_right", right)
record("full_total", total)
print(f"depth {depth_of(FULL)}, {len(leaves(FULL))} leaves, "
      f"{sum(1 for l in leaves(FULL) if l['note'].startswith('1 films'))} of them holding a single film")
print(f"gets {right} of {total} right — on the very films it was built from")

def leaf_counts(node, films=None):
    """Walk the tree with the films, so a leaf's contents are read not parsed."""
    if films is None:
        films = ALL
    if "leaf" in node:
        return [(node, films)]
    q = next(k for k in KEYS if SHORT[k] == node["q"])
    yes, no = split(films, q)
    return leaf_counts(node["yes"], yes) + leaf_counts(node["no"], no)

mixed = [(l, fs) for l, fs in leaf_counts(FULL) if 0 < sum(hit(f) for f in fs) < len(fs)]
for l, fs in mixed:
    print(f"leaf calls '{l['leaf']}'  holding {', '.join(fs)}")

record("n_mixed_leaves", len(mixed))
record("full_wrong", total - right)
record("stuck_films", sorted(f for _, fs in mixed for f in fs))

print(f"{'min films in a leaf':>20} {'depth':>6} {'leaves':>7} {'right of 24':>12}")
sweep = []
for m in range(1, 7):
    t = build(ALL, min_leaf=m)
    r, _ = score(t)
    sweep.append((m, depth_of(t), len(leaves(t)), r))
    print(f"{m:>20} {depth_of(t):>6} {len(leaves(t)):>7} {r:>12}")

record("sweep", [{"min_leaf": m, "depth": d, "leaves": lv, "right": r} for m, d, lv, r in sweep])
record("readable_min_leaf", 3)
READABLE = build(ALL, min_leaf=3)
record("readable_leaves", len(leaves(READABLE)))
record("readable_depth", depth_of(READABLE))
record("readable_right", score(READABLE)[0])
record("readable_tree", READABLE)

def plot2(ax):
    ms = [m for m, *_ in sweep]
    ax.plot(ms, [lv for *_, lv, _ in sweep], marker="o", ms=5, color="#9aa0aa", label="leaves in the tree")
    ax.plot(ms, [r for *_, r in sweep], marker="o", ms=5, color="#e2574c", label="films called right, of 24")
    ax.set_xlabel("Fewest films allowed in a leaf")
    ax.set_ylim(0, 26)
    ax.legend(frameon=False, fontsize=9)

save_fig("shrinking-the-tree", plot2, figsize=(7, 3.6))