Single Transferable Vote

stv.py

#!/usr/bin/env python
from math import floor
import sys


if len(sys.argv) < 2:
    print 'Usage: python stv.py [ballot filename]'
    exit()

DELIM = ','

seats = 3


class candidate:
    def __init__(self, name):
        self.name = name
        self.eligible = True
        self.elected = False
        self.score = 0.

    def __cmp__(self, other):
        return self.score - other.score

    def __repr__(self):
        return self.name


class voter:
    def __init__(self, votes):
        self.votes = votes
        self.used = 0
        self.value = 1.


with open(sys.argv[1]) as f:
    lines = f.readlines()
    candidates = [candidate(c) for c in lines[0].strip().split(DELIM)]

    # this needs to change
    votes = [voter([candidates[v.index(str(i))]
             for i in range(len(candidates)) if str(i) in v])
             for v in [line.strip().split(',')
             for line in lines[1:]]]

threshold = floor(1 + len(votes) / (seats + 1))

print '{} candidates:'.format(len(candidates))
for c in candidates:
    print '- {}'.format(c)
print ''
print '{} seats'.format(seats)
print '{} votes'.format(len(votes))
print ''
print 'Threshold: {:.2f}'.format(threshold)
print ''


def ended():
    return len(filter(lambda c: c.eligible, candidates)) == seats or \
           len(filter(lambda c: c.score >= threshold, candidates)) == seats


def show_scores():
    for c in candidates:
        if not c.eligible:
            continue

        status = ''

        if c.score > threshold and not c.elected:
            status = '(elected, surplus: {:.2f})'.format(c.score - threshold)
        if c.score == threshold or c.elected:
            status = '(elected)'

        print '{:<15}{:<8.2f}{}'.format(c, c.score, status)


def redistribute(c, value_factor=1.):
    for v in votes:
        if v.used < len(v.votes) and v.votes[v.used] == c:
            while v.used < len(v.votes) and \
                    (v.votes[v.used].score >= threshold or
                     not v.votes[v.used].eligible or v.votes[v.used] == c):
                v.used += 1

            if v.used < len(v.votes):
                v.value *= value_factor
                v.votes[v.used].score += v.value


round = 0

while not ended():
    print 'Round {}'.format(round + 1)
    print ''

    n = False

    for c in candidates:
        if c.score >= threshold and not c.elected:
            n = True
            c.elected = True
            if c.score > threshold:
                surplus_value = float(c.score - threshold) / c.score
                print \
                    '{} has {:.2f} surplus votes, which will be' \
                    ' redistributed with a value of {:.2f} per vote.'.format(
                        c, c.score - threshold, surplus_value)
                redistribute(c, surplus_value)

    if round == 0:
        print 'Adding most preferred votes'
        for v in votes:
            v.votes[0].score += 1.
    elif not n:
        left = filter(lambda c: c.eligible, candidates)
        elim = filter(lambda c: c.score == min(left).score, left)

        for c in elim:
            print '{} has been eliminated.'.format(c)
            if c.score > 0:
                print 'Their {:.2f} votes will be redistributed.'.format(
                    c.score)
            c.eligible = False
            redistribute(c)
            c.score = 0.

    print ''
    show_scores()
    print ''
    print ''

    round += 1

print 'Winners are {}'.format(', '.join(c.name for c in filter(
    lambda c: c.eligible and c.score >= threshold, candidates)))