stv.py

#!/usr/bin/env python3
import csv
from math import floor
import sys


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

seats = 5


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

    def __lt__(self, other):
        return self.score < other.score


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


with open(sys.argv[1], 'r') as f:
    reader = csv.reader(f)
    candidates = list(map(candidate, next(reader)))

    votes = []

    for row in reader:
        prefs = list(zip(row, candidates))
        prefs.sort(key=lambda t: t[0])
        votes.append(voter([t[1] for t in prefs if t[0] != '']))

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

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


def ended():
    return len([c for c in candidates if c.eligible]) == seats or \
           len([c for c in candidates if c.score >= threshold]) == seats


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

        status = ''

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

        print(f'{c.name:<15}{c.score:<8.2f}{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(f'Round {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(f'{c.name} has {(c.score - threshold):.2f} surplus votes, '
                      f'which will be redistributed with a value of '
                      f'{surplus_value:.2f} per vote.')
                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 = list(filter(lambda c: c.eligible, candidates))
        elim = filter(lambda c: c.score == min(left).score, left)

        for c in elim:
            print(f'{c.name} has been eliminated')
            if c.score > 0:
                print(f'Their {c.score:.2f} votes will be redistributed.')
            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))))