From Max Weinreich, a purge puzzle:

A town of 1,000 households has a strange law intended to prevent wealth-hoarding. On January 1 of every year, each household robs one other household, selected at random, moving all of that house’s money into their own house. The order in which the robberies take place is also random and is determined by a lottery. (Note that if House A robs House B first, and then C robs A, the houses of A and B would each be empty and C would have acquired the resources of both A and B.)

Two questions about this fateful day:

What is the probability that a house is not robbed over the course of the day?
Suppose that every house has the same amount of cash to begin with — say $100. Which position in the lottery has the most expected cash at the end of the day, and what is that amount?



In [1]:

    
from itertools import repeat, cycle, chain, count
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib as mpl
from tqdm import tqdm_notebook
from multiprocessing import Pool
import matplotlib.animation as animation
import numpy as np



In [2]:

    
def calculate_all_expecteds(houses):
    return list(map(calculate_expected, houses))



In [3]:

    
def calculate_expected(house):
    return sum([amt*odd for amt,odd in house.items()])



In [4]:

    
def possible_amounts(all_houses, ignore_index, num_houses):
    possibles = defaultdict(float)
    for j, house in enumerate(all_houses):
        if j != ignore_index:
            for amount, odd in house.items():
                possibles[amount] += odd

    return {a: o / (num_houses - 1) for a, o in possibles.items()}



In [5]:

    
def adjust_thief_odds(house, possibles):
    resulting_probs = []

    for h_amount, h_odd in house.items():
        for p_amount, p_odd in possibles.items():
            resulting_probs.append((h_amount + p_amount, h_odd * p_odd))

    new_house = defaultdict(float)
    for r_amount, r_odd in resulting_probs:
        new_house[r_amount] += r_odd

    return dict(new_house)



In [6]:

    
def adjust_victim_house(house, num_houses):
    total = 0
    new_house = defaultdict(float)
    for amount, odd in house.items():
        temp_odd = odd * (1 - (1/(num_houses - 1)))
        new_house[amount] = temp_odd
        total += temp_odd

    new_house[0] += 1 - total

    return dict(new_house)



In [7]:

    
def run_simulation(num_houses, initial_amount):
    all_houses = list(repeat({initial_amount: 1.0}, num_houses))

    for i, thief_house in enumerate(tqdm_notebook(all_houses, desc='calculation')):
        # calculate possible amounts that could be stolen, probability of each
        possibles = possible_amounts(all_houses, i, num_houses)

        # adjust probabilities of each amount in selected house
        all_houses[i] = adjust_thief_odds(thief_house, possibles)

        # add probability to each other house that their amount was stolen
        for j, victim_house in enumerate(all_houses):
            if j != i:
                all_houses[j] = adjust_victim_house(victim_house, num_houses)

        yield calculate_all_expecteds(all_houses)



In [71]:

    
def generate_evolution_figure(expecteds):
    num_houses = len(expecteds)
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    ax.set_xlim([0, num_houses])
    ax.set_ylim([0, max(expecteds[1])])
    ax.set_xlabel("House Index")
    ax.set_ylabel("Expected Amount [$]")
    ax.set_title("Expected Amount in House over time")
    
    bar = ax.bar(range(num_houses), expecteds[0], width=1, align='center')
    fit_line, = ax.plot([],[], 'r--', linewidth=3)
    
    return fig, bar, fit_line

def update_evolution_figure(frame_num,
                            fig,
                            bar,
                            fit_line,
                            norm,
                            expected_len,
                            expected_generator,
                            pbar):
    exp = next(expected_generator)
    
    for i, b in enumerate(bar):
        b.set_height(exp[i])
        b.set_facecolor(color=cm.ScalarMappable(norm=norm, cmap=cm.viridis).to_rgba(exp[i]))
    
    if frame_num >= expected_len:
        x = range(len(exp))
        p = np.poly1d(np.polyfit(x, exp, 2))
        fit_line.set_data(x, p(x))
    
    pbar.update(1)

def plot_expected_evolution(expecteds, output_file, fps, repeat_delay):
    fig, bar, fit_line = generate_evolution_figure(expecteds)
    norm = mpl.colors.Normalize(vmin=0, vmax=max(expecteds[1]))
    
    with tqdm_notebook(total=len(expecteds)+repeat_delay) as pbar:
        evo_ani = animation.FuncAnimation(fig,
                                          update_evolution_figure,
                                          fargs=(fig,
                                                 bar,
                                                 fit_line,
                                                 norm,
                                                 len(expecteds),
                                                 cycle(chain(expecteds,
                                                             repeat(expecteds[-1],
                                                                    repeat_delay))),
                                                 pbar),
                                          repeat=True,
                                          save_count=len(expecteds)+repeat_delay)
        evo_ani.save(output_file, writer="imagemagick", fps=fps)



In [72]:

    
NUM_HOUSES = 1000
INIT_AMOUNT = 100



In [73]:

    
expecteds = list(run_simulation(NUM_HOUSES, INIT_AMOUNT))



In [75]:

    
plot_expected_evolution(expecteds,'evolution.gif', 48, 100)
plt.close('all')