hacky implementation of hugo runoffs

hugo.py

#!/usr/bin/python3

import copy
import math
import pprint

candidates = [
    "Green Mars",
    "The Fifth Season",
    "What If",
    "The Dictator's Handbook",
    "Stories of Your Life and Others",
    "Pnin",
    "Three Body Problem",
]
raw_data = """Jeremy
The Fifth Season
Green Mars
The Dictator's Handbook
What If
Three Body Problem
Pnin

Alexis
What If
The Fifth Season
Three Body Problem
The Dictator's Handbook
Pnin
Green Mars

Tim
Three Body Problem
Pnin
The Fifth Season
Green Mars
What If
The Dictator's Handbook

Piotr
The Fifth Season
Pnin
The Dictator's Handbook
What If
Green Mars
Stories of Your Life and Others
Three Body Problem

Matthew
Pnin
Green Mars
The Dictator's Handbook
What If
Three Body Problem
The Fifth Season

Taylor
Pnin
Three Body Problem
Green Mars
What If
The Dictator's Handbook
The Fifth Season

Marty
Green Mars
Stories of Your Life and Others
What If
Pnin"""

raw_people_data = raw_data.split("\n\n")
votes_by_person = { raw_person_data.splitlines()[0] : raw_person_data.splitlines()[1:] for raw_person_data in raw_people_data }

counts_by_candidate = {}
for candidate in candidates:
    counts_by_candidate[candidate] = 0

# Validate that people actually voted for the candidates
for votes in votes_by_person.values():
    counts_by_candidate[votes[0]] += 1
    for vote in votes:
        assert vote in candidates

def sorted_candidate_counts(counts_by_candidate):
    return sorted(counts_by_candidate.items(), key=lambda candidate_count: candidate_count[1])

def runoff(votes_by_person, counts_by_candidate, candidate_to_remove=None, round_count=0):
    votes_by_person = copy.deepcopy(votes_by_person)
    counts_by_candidate = copy.deepcopy(counts_by_candidate)

    scc = sorted_candidate_counts(counts_by_candidate)

    best_candidate, best_candidate_count = scc[-1]
    if best_candidate_count > math.floor(len(votes_by_person) / 2):
        print("!"*35 + " Found a winner! " + "!"*35)
        print(best_candidate)
        return

    print("*"*35 + " Round {} ".format(round_count) + "*"*35)
    print("Removing candidate {}".format(candidate_to_remove))
    if candidate_to_remove:
        del counts_by_candidate[candidate_to_remove]

    for voter, votes in votes_by_person.items():
        # If this person voted for a removed candidate, transfer his vote!
        # This requires searching until we find a candidate they
        # voted for who is still in the running (there may be no
        # such vote).
        if len(votes) > 0 and votes[0] not in counts_by_candidate:
            original_vote = votes[0]
            votes.pop(0)
            while len(votes) > 0 and votes[0] not in counts_by_candidate:
                votes.pop(0)
            if len(votes) > 0 and votes[0] in counts_by_candidate:
                print("Transferring {}'s vote from {} to {}".format(voter, original_vote, votes[0]))
                counts_by_candidate[votes[0]] += 1

    scc = sorted_candidate_counts(counts_by_candidate)

    print("Votes after transferring votes in round {}".format(round_count))
    pprint.pprint(scc)

    fewest_votes_count = scc[0][1]
    losers = [ candidate_count[0] for candidate_count in scc if candidate_count[1] == fewest_votes_count ]
    print("Losers {}".format(losers))
    for loser in losers:
        runoff(votes_by_person, counts_by_candidate, candidate_to_remove=loser, round_count=round_count+1)

print("Initial votes")
pprint.pprint(votes_by_person)
runoff(votes_by_person, counts_by_candidate)