Cleaned up solution to Gerrymandering (HackMIT Puzzle Hunt 2018)

gerrymandering.py

import copy
import random

import numpy as np
import pandas as pd
from scipy.stats import norm

pa, pb = pd.read_json('voters.json').voters_by_block
population_a, population_b = np.array(pa), np.array(pb)
population = population_a + population_b

def win_probability(na, nb):
    mean = .6 * na + .4 * nb
    variance = (na + nb) * .4 * .6
    z = ((na + nb) / 2.0 - mean) / (variance ** 0.5)
    return 1 - norm.cdf(z)

def expected_districts_won(soln):
    ret = 0
    for district in soln:
        na = np.sum(population_a[district])
        nb = np.sum(population_b[district])
        ret += win_probability(na, nb)
    return ret

def district_population_imbalance(soln):
    sizes = [np.sum(population[district]) for district in soln]
    return np.var(sizes) * len(sizes)

def expected_efficiency_gap(soln):
    # This is approximate only
    # Calculating it exactly requires some messy integrals
    total = 0.2 * sum(population_a) - 0.2 * sum(population_b)
    for district in soln:
        na = np.sum(population_a[district])
        nb = np.sum(population_b[district])
        p = win_probability(na, nb)
        ea = 0.6 * na + 0.4 * nb
        eb = 0.6 * nb + 0.4 * na
        used = p * eb - (1 - p) * ea # used for a - used for b
        total -= used
    return total / sum(population)

def neighbors(idx):
    x, y = divmod(idx, 10)
    ret = []
    if x: ret.append(idx - 10)
    if y: ret.append(idx - 1)
    if y < 9: ret.append(idx + 1)
    if x < 9: ret.append(idx + 10)
    random.shuffle(ret)
    return ret

def connected(d):
    if not len(d):
        # We can't have empty districts...
        return False
    vis = { city: False for city in d }
    def dfs(city):
        if vis[city]:
            return
        vis[city] = True
        for n in neighbors(city):
            if n in d:
                dfs(n)
    dfs(d[0])
    return all(vis.values())

def valid(soln):
    return all(connected(district) for district in soln)

def random_neighbor(soln):
    dnum = {}
    for i, district in enumerate(soln):
        for city in district:
            dnum[city] = i
    while True:
        city = random.randint(0, 99)
        for n in neighbors(city):
            if dnum[city] != dnum[n]:
                # Attempt to switch city's district to n's district
                soln2 = copy.deepcopy(soln)
                soln2[dnum[city]].remove(city)
                soln2[dnum[n]].append(city)
                if valid(soln2):
                    return soln2

# Simulated Annealing
def P(e, e2, t):
    if e2 < e:
        return 1
    return np.exp(-(e2 - e) / t)

def accept(e, e2, temp):
    return np.random.rand() < P(e, e2, temp)

def temperature(completed):
    # Exponential temperature function
    return 100 * np.exp(-10 * completed)

def properties(soln):
    return (expected_districts_won(soln),
        district_population_imbalance(soln),
        expected_efficiency_gap(soln))

def energy(soln):
    X_MIN = 10 * 1.001
    Y_MAX = 164e10 * .999
    Z_MIN = -0.074 * .999

    ret = 0
    x, y, z = properties(soln)
    if x < X_MIN:
        ret += X_MIN - x
    if y > Y_MAX:
        ret += (y / Y_MAX - 1) * 3
    if z < Z_MIN:
        ret += (Z_MIN - z) * 2
    return ret

def main():
    ITERATIONS = 10000
    soln = [list(range(x, x + 5)) for x in range(0, 100, 5)]
    e = energy(soln)
    for iteration in range(ITERATIONS):
        temp = temperature(iteration / ITERATIONS)
        if iteration % 100 == 99:
            print('Iteration {}, energy = {:.2f}, temp = {:.2f}'.format(iteration + 1, e, temp))
        soln2 = random_neighbor(soln)
        e2 = energy(soln2)
        if accept(e, e2, temp):
            e, soln = e2, soln2
            if e2 < 1e-8:
                break
    print(properties(soln))
    print(soln)

if __name__ == '__main__':
    main()