Updated geryrmandering code

gerry.py

import z3
import pulp


def validate(districts, pro):
    if not districts:
        raise ValueError("Must have at least one district")
    for v in districts:
        if v <= 0:
            raise ValueError("Invalid district size %r" % (v,))
    population = sum(districts)
    if pro <= 0:
        raise ValueError("Invalid pro population %r" % (pro,))
    if pro > population:
        raise ValueError("Can't have more people pro that population")
    if pro == population:
        raise ValueError("Can't have everyone pro")


def gerrymanderz3(districts, pro):
    districts = list(districts)
    validate(districts, pro)

    solver = z3.Optimize()

    populations = [z3.Int('Pop%d' % (i,)) for i in range(len(districts))]
    wins = []
    for i, (p, d) in enumerate(zip(populations, districts)):
        solver.add(p >= 0)
        solver.add(p <= d)
        win = z3.Int('Win%d' % (i,))
        solver.add(z3.Implies(
            p * 2 > d, win == 1
        ))
        solver.add(z3.Implies(
            p * 2 <= d, win == 0
        ))
        wins.append(win)
    solver.add(sum(populations) == pro)
    wincount = z3.Int('wincount')
    solver.add(wincount == sum(wins))
    assert solver.check() == z3.sat

    lo = 0
    hi = len(districts) + 1
    # Invariant: There is a partition with at least lo districts majority pro
    # Invariant: There is no partition with at least hi districts majority pro
    while lo + 1 < hi:
        mid = (lo + hi) // 2
        solver.push()
        solver.add(wincount >= mid)
        check = solver.check()
        if check == z3.unsat:
            hi = mid
        else:
            lo = mid
        solver.pop()
    solver.add(wincount >= lo)
    res = solver.check()
    assert res == z3.sat, res
    m = solver.model()
    return [m[p] for p in populations], lo


def gerrymanderlp(districts, pro):
    validate(districts, pro)
    problem = pulp.LpProblem('gerrymander', pulp.LpMaximize)

    votes = [
        pulp.LpVariable(
            'Pop%d' % (i,), cat='Integer',
            lowBound=0, upBound=v
        )
        for i, v in enumerate(districts)
    ]
    wins = []
    for i, (population, count) in enumerate(zip(districts, votes)):
        win = pulp.LpVariable(
            'Win%d' % (i,), cat='Integer', lowBound=0, upBound=1)
        wins.append(win)
        loss_ceiling = population // 2
        if loss_ceiling > 0:
            loss = pulp.LpVariable(
                'Loss%d' % (i,), cat='Integer', lowBound=0,
                upBound=loss_ceiling)
            problem.add(loss * (1.0 / loss_ceiling) >= win)
        else:
            loss = 0
        problem.add(loss + win <= count)

    assert len(wins) == len(votes) == len(districts)
    problem.add(sum(votes) == pro)

    problem.objective = sum(wins)
    problem.solve()
    votes_per_district = [int(p.varValue) for p in votes]
    votes_per_district.sort(reverse=True)
    return (
        votes_per_district,
        int(sum(w.varValue for w in wins))
    )


def gerrymandergreedy(districts, pro):
    districts = sorted(districts)
    wins = 0
    results = []
    for d in districts:
        win = (d // 2) + 1
        if win <= pro:
            pro -= win
            results.append(win)
            wins += 1
        else:
            pro = 0
            results.append(pro)
    return results, wins

test_gerry.py

from gerry import gerrymanderz3, gerrymanderlp, gerrymandergreedy
from hypothesis import given, assume, example
from hypothesis import strategies as st
import pytest


@st.composite
def problem(draw):
    districts = draw(
        st.lists(
            st.integers(min_value=1, max_value=100),
            min_size=1, max_size=5)
    )
    population = sum(districts)
    assert population >= 1
    assume(population >= 2)
    return districts, draw(st.integers(1, population - 1))


implementations = [gerrymanderz3, gerrymanderlp, gerrymandergreedy]

gerrytest = pytest.mark.parametrize(
    'gerrymander', implementations)


@gerrytest
def test_can_gerrymander_majority_to_total(gerrymander):
    _, count = gerrymander([3] * 3, 6)
    assert count == 3


@gerrytest
def test_can_arrange_a_win_for_a_minority(gerrymander):
    _, count = gerrymander([3] * 3, 2)
    assert count == 1


@gerrytest
@given(problem())
def test_handles_arbitrary_data(gerrymander, problem):
    gerrymander(*problem)


@gerrytest
@given(problem())
def test_increasing_population_increases_maximum(gerrymander, problem):
    districts, pro = problem
    assume(pro + 1 < sum(districts))
    _, c1 = gerrymander(districts, pro)
    _, c2 = gerrymander(districts, pro + 1)
    assert c1 <= c2


@example(([2, 2], 2))
@example(([1, 1], 1))
@given(problem())
def test_counts_agree(problem):
    cs = [g(*problem) for g in implementations]
    assert len({c for _, c in cs}) == 1