coupon_problem.py

# Written by Claude (Anthropic AI)
# Rewritten by Za3k because it was bad style.

from decimal import Decimal, getcontext
from math import comb
from itertools import product
from collections import defaultdict

getcontext().prec = 100
def coupon_collector(probs):
    book_probs = defaultdict(int)
    for p in probs:
        book_probs[p] += 1

    # Calculate using inclusion-exclusion
    # Group by prob only (to take advantage of repeated structure and avoid 2^40 subsets)
    group_probs = list(book_probs.keys())
    group_sizes = [book_probs[p] for p in group_probs]

    E_T = Decimal(0)
    ranges = [range(size + 1) for size in group_sizes]
    for state in product(*ranges):
        if all(n == 0 for n in state):
            continue
        
        total_in_state = sum(state)
        
        multiplicity = 1
        for i, count in enumerate(state):
            multiplicity *= comb(group_sizes[i], count)
        
        p_sum = Decimal(0)
        for i, count in enumerate(state):
            prob = group_probs[i]
            p_sum += Decimal(count) * prob
        
        sign = (-1) ** (total_in_state + 1)
        contribution = Decimal(multiplicity * sign) / p_sum
        E_T += contribution
    return E_T

ENCHANTS = """Aqua Affinity I|Bane of Arthropods V|Blast Protection IV|Breach IV|Channeling I|Curse of Binding I|Curse of Vanishing I|Depth Strider III|Density V|Efficiency V|Feather Falling IV|Fire Aspect II|Fire Protection IV|Flame I|Fortune III|Frost Walker II|Impaling V|Infinity I|Knockback II|Looting III|Loyalty III|Luck of the Sea III|Lunge III|Lure III|Mending I|Multishot I|Piercing IV|Power V|Projectile Protection IV|Protection IV|Punch II|Quick Charge III|Respiration III|Riptide III|Sharpness V|Silk Touch I|Smite V|Sweeping Edge III|Thorns III|Unbreaking III""".split("|")
enchantments = []
for line in ENCHANTS:
    parts = line.split()
    name = ' '.join(parts[:-1])
    level_map = {'I': 1, 'II': 2, 'III': 3, 'IV': 4, 'V': 5}
    max_level = level_map[parts[-1]]
    # turns out "is_treasure" wouldn't be useful
    
    enchantments.append((name, max_level)) 

n_enchants = len(enchantments)
print(f"Total enchantments: {n_enchants}")

ps1 = []
ps2 = []
ps3 = []
# Calculate probability for each book under various requirements
for name, max_level in enchantments:
    book_offer_prob = Decimal(2) / Decimal(3)

    # Step 1: Pick this enchantment
    p_enchant = Decimal(1) / Decimal(n_enchants)
    
    # Step 2: Pick max level
    p_level = Decimal(1) / Decimal(max_level)

    # Step 2: Pick max level
    p_level = Decimal(1) / Decimal(max_level)
    
    # Step 3: Determine range
    # There's a cap at 64, and treasure enchantments double... but that turns out not to matter.
    random_range = 5 + max_level * 10
    p_price = Decimal(1) / Decimal(random_range)
    
    # Combined probability
    ps1.append(book_offer_prob * p_enchant)
    ps2.append(book_offer_prob * p_enchant * p_level)
    ps3.append(book_offer_prob * p_enchant * p_level * p_price)

print()
E_T = coupon_collector(ps1)
print(f"Expected villagers at any level: {float(E_T):.1f}")
E_T = coupon_collector(ps2)
print(f"Expected villagers at max enchant level: {float(E_T):.1f}")
E_T = coupon_collector(ps3)
print(f"Expected villagers at optimal prices: {float(E_T):.1f}")