runoff: reuse IRV Numba tally helpers for head-to-head round

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
         args: [--maxkb=2048]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.14
+    rev: v0.15.16
     hooks:
       # Quote --select=...: in YAML flow style, commas inside [...] split list items (E902 on Windows).
       - id: ruff

elsim/methods/runoff.py

@@ -1,7 +1,14 @@
 import numpy as np
 
-from elsim.methods._common import (_all_indices, _get_tiebreak, _no_tiebreak,
-                                   _order_tiebreak_keep, _random_tiebreak)
+from elsim.methods._common import (
+    _all_indices,
+    _get_tiebreak,
+    _inc_rank_idx,
+    _no_tiebreak,
+    _order_tiebreak_keep,
+    _random_tiebreak,
+    _tally_at_rank_idx,
+)
 
 _tiebreak_map = {'order': _order_tiebreak_keep,
                  'random': _random_tiebreak,
@@ -93,19 +100,19 @@ def runoff(election, tiebreaker=None):
     if None in finalists:
         return None
 
-    # TODO: Can this be vectorized or numbafied?
     finalist_0 = finalists[0]
     finalist_1 = finalists[1]
 
-    finalist_0_tally = 0
-    finalist_1_tally = 0
-
-    for ballot in election:
-        ballot_list = ballot.tolist()
-        if ballot_list.index(finalist_0) < ballot_list.index(finalist_1):
-            finalist_0_tally += 1
-        else:
-            finalist_1_tally += 1
+    n_cands = election.shape[1]
+    cand_tallies = np.empty(n_cands, dtype=np.uint)
+    voter_top_rank_idx = np.zeros(n_voters, dtype=np.uint8)
+    eliminated_mask = np.ones(n_cands, dtype=bool)
+    eliminated_mask[finalist_0] = False
+    eliminated_mask[finalist_1] = False
+    _inc_rank_idx(election, voter_top_rank_idx, eliminated_mask)
+    _tally_at_rank_idx(cand_tallies, election, voter_top_rank_idx)
+    finalist_0_tally = cand_tallies[finalist_0]
+    finalist_1_tally = cand_tallies[finalist_1]
 
     assert finalist_0_tally + finalist_1_tally == n_voters
 

tests/test_runoff.py

@@ -6,6 +6,7 @@
 from hypothesis.strategies import integers, lists, permutations
 
 from elsim.methods import runoff
+from elsim.methods._common import _inc_rank_idx, _tally_at_rank_idx
 
 
 def collect_random_results(method, election):
@@ -22,6 +23,41 @@ def collect_random_results(method, election):
     return winners
 
 
+def test_docstring_example():
+    """Contingent vote example from runoff() docstring."""
+    A, B, C = 0, 1, 2
+    election = np.array([[A, C, B],
+                         [A, C, B],
+                         [B, C, A],
+                         [B, C, A],
+                         [C, A, B]])
+    assert runoff(election) == A
+
+
+def test_head_to_head_tally_via_irv_helpers():
+    """
+    Second-round tally skips eliminated candidates then counts top preference
+    among the two finalists (same path as IRV).
+    """
+    A, B, C = 0, 1, 2
+    election = np.array([[A, C, B],
+                         [A, C, B],
+                         [B, C, A],
+                         [B, C, A],
+                         [C, A, B]])
+    n_voters, n_cands = election.shape
+    voter_top_rank_idx = np.zeros(n_voters, dtype=np.uint8)
+    cand_tallies = np.empty(n_cands, dtype=np.uint)
+    eliminated_mask = np.ones(n_cands, dtype=bool)
+    eliminated_mask[A] = False
+    eliminated_mask[B] = False
+    _inc_rank_idx(election, voter_top_rank_idx, eliminated_mask)
+    _tally_at_rank_idx(cand_tallies, election, voter_top_rank_idx)
+    assert cand_tallies[A] == 3
+    assert cand_tallies[B] == 2
+    assert cand_tallies[C] == 0
+
+
 def test_one_round():
     # 60% majority, tie between others
     election = np.array([[2, 0, 1],