Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Updated geryrmandering code
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.