Process ranked-choice votes via the topological sort of head-to-heads

rcv_toposort.py

# Copyright 2024 Google LLC.
# SPDX-License-Identifier: Apache-2.0

from cyclic_toposort import cyclic_toposort, acyclic_toposort
from collections import defaultdict
import re

votes = """
Shane: 1 ~> 3a > (3b ~> 5) > (2 ~> 4)
Manish: (1 ~> 3b) > 3a > 5 > 2 >> 4
Zibi: 3a ~> 3b > 5 ~> 2 ~>> 4 > 1
Elango: 3a ~> 3b > 4 ~>> 2 >> (1 = 5)
Robert: 1 > 3a ~>> 3b > 2 >> 4 = 5
"""

def parse_line_to_edges(line):
	"""
	Takes a voting line and yields items of the following form:

	(a, b, w)

	where `a` is the left node, `b` is the right node, and `w` is the weight.
	There is an initial edge with `a = None`.

	For example: "Shane: 1 ~>> 3 ~> 4 > 2" results in the following:

	(None, "1", 0)
	("1", "3", 3)
	("3", "4", 1)
	("4", "2", 2)
	"""
	tokens = iter(line.split(" "))
	name = next(tokens)
	prev_item = re.sub(r"[\(\)]", "", next(tokens))
	yield (None, prev_item, 0)
	edges = []
	while True:
		operator = next(tokens, None)
		if operator is None:
			break
		next_item = re.sub(r"[\(\)]", "", next(tokens))
		print(name, prev_item, operator, next_item)
		if operator == "~>":
			yield (prev_item, next_item, 1)
		elif operator == ">":
			yield (prev_item, next_item, 2)
		elif operator == "~>>":
			yield (prev_item, next_item, 3)
		elif operator == ">>":
			yield (prev_item, next_item, 5)
		elif operator == "=":
			yield (prev_item, next_item, 0)
		else:
			raise ValueError("Unknown operator: %s" % operator)
		prev_item = next_item
	edges

def format_set(s):
	"""
	Formats a set of strings. If there is a single element, that element is
	returned; otherwise, they are formatted as "(a = b = c)".
	"""
	if len(s) == 1:
		return next(iter(s))
	else:
		return "(%s)" % " = ".join(sorted(s))

def format_sets(ss):
	"""
	Formats an ordered iterator of sets of strings.

	For example: [{"a"}, {"b", "c"}] is formatted as: "a > (b = c)"
	"""
	return " > ".join(format_set(s) for s in ss)

# Walk through the votes to calculate the head-to-head weights.
total_weights = defaultdict(lambda: defaultdict(int))
for line in votes.split("\n"):
	if not line:
		continue
	current_weight = 0
	local_weights = dict()
	for (a, b, w) in parse_line_to_edges(line):
		current_weight += w
		local_weights[b] = current_weight
		for (x, w1) in local_weights.items():
			total_weights[x][b] += (current_weight - w1)

all_nodes = list(sorted(set(total_weights.keys())))

# We'll evaluate two different ways: total number of head-to-heads won,
# and a topological sort that minimizes dropped edges.

head2heads = defaultdict(int)
edges_for_toposort = []

for a in all_nodes:
	for b in all_nodes:
		net = total_weights[a][b] - total_weights[b][a]
		if net > 0:
			head2heads[a] += 1
			head2heads[b] -= 1
			edges_for_toposort.append((a, b))
			print(a, ">", b, "by a weight of", net)

stack_rank = defaultdict(set)
for (x, w) in head2heads.items():
	stack_rank[-w].add(x)

# Print out the results of both algorithms. Usually, they are the same.
print("Stack rank:", format_sets(x[1] for x in sorted(stack_rank.items())))
print("Toposort:", format_sets(cyclic_toposort(set(edges_for_toposort))[0]))