Greedy search — why the best step is not the best path

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?',
}
KEYS = list(QUESTIONS)
SHORT = {'star': 'A star?', 'summer': 'Opened in summer?', 'sequel': 'A sequel?',
         'wide': '3,000+ screens?', 'budget': 'Over $100m?'}

#                          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)
answers = lambda f: FILMS[f][:len(KEYS)]
hit = lambda f: FILMS[f][-1]
record('n_films', len(ALL))
record('n_questions', len(KEYS))
print(f'{len(ALL)} films, {sum(hit(f) for f in ALL)} of them hits')

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) ** 2

def gain(films, q):
    a, b = split(films, q)
    return gini(films) - (len(a) * gini(a) + len(b) * gini(b)) / len(films)

record('root_gini', round(gini(ALL), 3))
for q in sorted(KEYS, key=lambda q: -gain(ALL, q)):
    print(f'{QUESTIONS[q]:<42} gain {gain(ALL, q):.3f}')

greedy_root = max(KEYS, key=lambda q: gain(ALL, q))
g_yes, g_no = split(ALL, greedy_root)
greedy_l = max([q for q in KEYS if q != greedy_root], key=lambda q: gain(g_yes, q))
greedy_r = max([q for q in KEYS if q != greedy_root], key=lambda q: gain(g_no, q))

def leaves_of(root, ql, qr):
    yes, no = split(ALL, root)
    out = []
    for side, q in ((yes, ql), (no, qr)):
        if q is None:
            out.append(side)
        else:
            a, b = split(side, q)
            out += [a, b]
    return out

def right(films):
    return max(sum(hit(f) for f in films), len(films) - sum(hit(f) for f in films))

score = lambda root, ql, qr: sum(right(l) for l in leaves_of(root, ql, qr))
total_gini = lambda root, ql, qr: sum(len(l) * gini(l) for l in leaves_of(root, ql, qr)) / len(ALL)

def as_tree(root, ql, qr):
    yes, no = split(ALL, root)
    def side(films, q):
        if q is None:
            return leaf(films)
        a, b = split(films, q)
        return {'q': SHORT[q], 'note': f'{len(films)} films - {sum(hit(f) for f in films)} hit',
                'yes': leaf(a), 'no': leaf(b)}
    def leaf(films):
        return {'leaf': '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'}
    return {'q': SHORT[root], 'note': f'{len(ALL)} films - {sum(hit(f) for f in ALL)} hit',
            'yes': side(yes, ql), 'no': side(no, qr)}

greedy_score = score(greedy_root, greedy_l, greedy_r)
greedy_gini = total_gini(greedy_root, greedy_l, greedy_r)
record('greedy_question', QUESTIONS[greedy_root])
record('greedy_gain', round(gain(ALL, greedy_root), 3))
record('greedy_right', greedy_score)
record('greedy_gini', round(greedy_gini, 3))
record('greedy_tree', as_tree(greedy_root, greedy_l, greedy_r))
print(f'greedy: {QUESTIONS[greedy_root]} -> {greedy_score}/{len(ALL)}, total gini {greedy_gini:.4f}')

trees = []
for root in KEYS:
    others = [q for q in KEYS if q != root]
    for ql in [None] + others:
        for qr in [None] + others:
            trees.append({'root': root, 'yes': ql, 'no': qr,
                          'right': score(root, ql, qr), 'gini': total_gini(root, ql, qr)})

record('n_trees', len(trees))
record('count_columns', ['', 'Choices'])
record('count_rows', [
    ['The first question', f'{len(KEYS)}'],
    ['On the yes side: stop, or one of the rest', f'{len(KEYS)}'],
    ['On the no side: stop, or one of the rest', f'{len(KEYS)}'],
    ['Trees in total', f'{len(KEYS)} x {len(KEYS)} x {len(KEYS)} = {len(trees)}'],
])
print(f'{len(trees)} trees')

by_gini = sorted(trees, key=lambda t: t['gini'])
lower = [t for t in by_gini if t['gini'] < greedy_gini - 1e-9]
best_g = by_gini[0]

record('n_lower_gini', len(lower))
record('greedy_gini_rank', len(lower) + 1)
record('best_gini', round(best_g['gini'], 3))
record('best_gini_question', QUESTIONS[best_g['root']])

def plot_gini(ax):
    vals = [t['gini'] for t in trees]
    ax.hist(vals, bins=28, color='#c9c9d1')
    ax.axvline(greedy_gini, color='#e2574c', lw=2.2)
    ax.axvline(best_g['gini'], color='#2f9e6e', lw=2.2, ls='--')
    ax.annotate(f"greedy  {greedy_gini:.3f}", (greedy_gini, ax.get_ylim()[1] * 0.92),
                textcoords='offset points', xytext=(7, 0), fontsize=9, color='#e2574c')
    ax.annotate(f"best  {best_g['gini']:.3f}", (best_g['gini'], ax.get_ylim()[1] * 0.72),
                textcoords='offset points', xytext=(-8, 0), ha='right', fontsize=9, color='#2f9e6e')
    ax.set_xlabel('Total impurity left in the finished tree (lower is better)')
    ax.set_ylabel('Trees')
    ax.grid(axis='x', visible=False)

save_fig('no-global-cost', plot_gini, figsize=(7, 3.6))
print(f"{len(lower)} of {len(trees)} trees have a lower total gini than greedy's {greedy_gini:.4f}")

best_r = max(trees, key=lambda t: t['right'])
better = [t for t in trees if t['right'] > greedy_score]

record('n_better', len(better))
record('best_right', best_r['right'])
record('best_question', QUESTIONS[best_r['root']])
record('best_gain_of_root', f"{gain(ALL, best_r['root']):.3f}")
record('best_tree', as_tree(best_r['root'], best_r['yes'], best_r['no']))

record('compare_columns', ['', 'First question', 'Films right', 'Total impurity'])
record('compare_rows', [
    ['The tree greedy builds', QUESTIONS[greedy_root], f'{greedy_score} of {len(ALL)}', f'{greedy_gini:.3f}'],
    ['The best two-question tree', QUESTIONS[best_r['root']], f"{best_r['right']} of {len(ALL)}", f"{best_r['gini']:.3f}"],
])

def plot_right(ax):
    vals = sorted({t['right'] for t in trees}, reverse=True)
    counts = [sum(1 for t in trees if t['right'] == v) for v in vals]
    colors = ['#2f9e6e' if v == best_r['right'] else '#e2574c' if v == greedy_score else '#c9c9d1'
              for v in vals]
    ax.bar([str(v) for v in vals], counts, color=colors)
    for i, (v, c) in enumerate(zip(vals, counts)):
        if v in (best_r['right'], greedy_score):
            ax.text(i, c + 1.2, 'best' if v == best_r['right'] else 'greedy', ha='center',
                    fontsize=9, color='#2f9e6e' if v == best_r['right'] else '#e2574c')
    ax.set_xlabel(f'Films right, out of {len(ALL)}')
    ax.set_ylabel('Trees')
    ax.grid(axis='x', visible=False)

save_fig('greedy-against-every-tree', plot_right, figsize=(7, 3.6))
print(f"greedy {greedy_score}/{len(ALL)}; best {best_r['right']}/{len(ALL)}; "
      f"{len(better)} tree(s) beat greedy")
print(f"best tree's first question: {QUESTIONS[best_r['root']]} — gain {gain(ALL, best_r['root']):.3f}")