An implementation of STV with the Droop quota, using the Gregory method and an ad-hoc tie breaking system

stv.py

from copy import copy
from collections import defaultdict
from fractions import Fraction

def stv(candidates, ballots, seats=1):
    if seats > len(candidates):
        raise ValueError('Not enough candidates for election')

    rounds = {}

    quota = 1 + Fraction(len(ballots), seats + 1)

    winners = set()
    remaining_candidates = set(candidates)

    ballot_map = defaultdict(lambda: 0)
    for ballot in ballots:
        ballot_map[tuple(ballot)] += 1

    round_number = 0
    while len(winners) < seats:
        round_number += 1
        if len(remaining_candidates) <= (seats - len(winners)):
            return {'winners': list(winners | remaining_candidates),
                    'quota': quota,
                    'rounds': rounds}
        # drop empty ballots
        ballot_map.pop((), 0)
        tally = {candidate: sum(weight for ballot, weight in ballot_map.items() if ballot[0] == candidate)
                   for candidate in remaining_candidates}
        rounds[round_number] = {'tally': tally}
        round_winners = {candidate for candidate, count in tally.items() if count >= quota}
        if round_winners:
            rounds[round_number]['winners'] = list(round_winners)
            winners |= round_winners
            for round_winner in round_winners:
                # eliminate ballots just for this winner
                ballot_map.pop((round_winner,), 0)
                overspill = tally[round_winner] - quota
                transfer_fraction = Fraction(overspill, tally[round_winner])
                old_ballot_map = copy(ballot_map)
                ballot_map.clear()
                remaining_candidates.remove(round_winner)
                for ballot, weight in old_ballot_map.items():
                    if ballot[0] == round_winner:
                        ballot_map[tuple(x for x in ballot[1:] if x in remaining_candidates)] += weight * transfer_fraction
                    else:
                        ballot_map[tuple(x for x in ballot if x in remaining_candidates)] += weight
        else:
            min_votes = min(count for candidate, count in tally.items())
            losers = {candidate for candidate, count in tally.items() if count == min_votes}
            rounds[round_number]['tied_losers'] = list(losers)
            drop = min(losers, key=lambda c: sum(weight for ballot, weight in ballot_map.items() if c in ballot))
            rounds[round_number]['loser'] = drop
            old_ballot_map = copy(ballot_map)
            ballot_map.clear()
            remaining_candidates.remove(drop)
            for ballot, weight in old_ballot_map.items():
                ballot_map[tuple(x for x in ballot if x in remaining_candidates)] += weight
    return {'winners': list(winners),
            'quota': quota,
            'rounds': rounds}

if __name__ == '__main__':
    import yaml
    with open('ballot2.yaml', 'r') as f:
        data = yaml.load(f)
    result = stv(candidates=data['candidates'].keys(),
                 ballots=data['votes'],
                 seats=data['seats'])
    for candidate, name in data['candidates'].items():
        print(' {} {}'.format('*' if candidate in result['winners'] else ' ', name))