fullrigger.py

import pulp
from itertools import permutations

SOLVER = pulp.solvers.COIN()


def lpsum(variables):
    result = pulp.LpAffineExpression()
    for v in variables:
        result.addInPlace(v)
    return result


def build_election(
    n_candidates, n_votes, additive_winners=(),
    condorcet_winner=None, irv_dropout_order=None,
):
    candidates = range(n_candidates)

    variables = [
        (p, pulp.LpVariable(
            'V_%s' % ('_'.join(map(str, p)),),
            lowBound=0, upBound=n_votes, cat='Integer'))
        for i, p in enumerate(permutations(candidates))
    ]

    non_zero_variables = [
        pulp.LpVariable('I%d' % (i,), cat='Binary')
        for i in range(len(variables))
    ]

    problem = pulp.LpProblem(sense=pulp.LpMinimize)

    problem.addConstraint(lpsum([v for _, v in variables]) == n_votes)

    for a, (_, v) in zip(non_zero_variables, variables):
        problem.addConstraint(a * n_votes >= v)

    problem.objective = lpsum(non_zero_variables)

    for scoring_function, winners in additive_winners:
        scores = {
            c: lpsum([scoring_function(p, c) * v for p, v in variables])
            for c in candidates
        }
        losers = set(candidates) - set(winners)
        for u, v in zip(winners, winners[1:]):
            problem.addConstraint(scores[u] >= scores[v] + 1)
        for u in winners:
            for v in losers:
                problem.addConstraint(scores[u] >= scores[v] + 1)

    if condorcet_winner is not None:
        victories = {
            (i, j): lpsum(
                v for p, v in variables if p.index(i) < p.index(j)
            )
            for i in candidates
            for j in candidates
        }
        for c in candidates:
            if c != condorcet_winner:
                problem.addConstraint(
                    victories[(condorcet_winner, c)] >=
                    victories[(c, condorcet_winner)] + 1
                )

    if irv_dropout_order is not None:
        remaining_candidates = set(candidates)
        for i in irv_dropout_order:
            if len(remaining_candidates) <= 1:
                break
            assert i in remaining_candidates
            allocations = {j: [] for j in remaining_candidates}
            for p, v in variables:
                for c in p:
                    if c in remaining_candidates:
                        allocations[c].append(v)
                        break
            loser_allocations = sum(allocations.pop(i))
            remaining_candidates.remove(i)
            for vs in allocations.values():
                problem.addConstraint(loser_allocations + 1 <= sum(vs))

    result = problem.solve(SOLVER)
    if pulp.LpStatus[result] == 'Optimal':
        return [
            (p, int(v.value())) for p, v in variables if v.value() > 0
        ]
    else:
        raise ValueError("No such election")