From 8732fbc43d17b02215e32d25827433430a56bebf Mon Sep 17 00:00:00 2001
From: Jack Wines <jack@winesj.com>
Date: Sun, 9 Jul 2023 01:48:26 -0700
Subject: [PATCH] InstantRunoff cleanup

---
 src/InstantRunoff.hs | 57 +++++++++++++++++++++++++++-----------------
 1 file changed, 35 insertions(+), 22 deletions(-)

diff --git a/src/InstantRunoff.hs b/src/InstantRunoff.hs
index b611779..f8b0a64 100644
--- a/src/InstantRunoff.hs
+++ b/src/InstantRunoff.hs
@@ -4,61 +4,74 @@ import qualified Data.List.NonEmpty as LN
 import qualified Data.Map.Strict as M
 import qualified Data.Set as S
 import qualified Data.Maybe as M
-import qualified Data.Ord as O
+import qualified Data.Foldable as F
 
 solve :: forall a. Ord a => [a] ->  LN.NonEmpty (LN.NonEmpty a) -> [S.Set a]
-solve candidates votes = (reverse solved) ++ if (S.null notVotedForCandidates) then [] else [notVotedForCandidates]
+solve candidates votes = reverse . (consIfAny notVotedForCandidates) . rank $ votes
   where
 
+    consIfAny :: S.Set a -> [S.Set a] -> [S.Set a]
+    consIfAny [] xs = xs
+    consIfAny x xs = x : xs
+
     -- candidates not on any submitted ballot
     notVotedForCandidates :: S.Set a
-    notVotedForCandidates = S.difference (S.fromList candidates) . S.unions $ solved
+    notVotedForCandidates =
+      S.difference (S.fromList candidates)
+      -- all candidates on any submited ballot
+      . S.unions . LN.map (S.fromList . LN.toList) $ votes
 
-    solved = rank votes
 
 rank :: forall a. Ord a => LN.NonEmpty (LN.NonEmpty a) -> [S.Set a]
-rank votes = maybe (L.singleton . M.keysSet $ voteCounts) ((losers :) . rank) . filterVotes (`S.notMember` losers) $ votes
+rank votes =
+  maybe [losers] ((losers :) . rank)
+  . filterVotes (`S.notMember` losers) $ votes
   where
 
+    -- losers in terms of the first preference on ballot
     firstChoiceLosers :: S.Set a
     firstChoiceLosers = M.keysSet . M.filter ((==) . L.minimum . M.elems $ voteCounts) $ voteCounts
 
     -- in the event of multiple first choice losers,
     -- count how many ballots prefer one choice over another
     losers :: S.Set a
-    losers = S.fromList . smallestSubset preferred . S.toList $ firstChoiceLosers
+    losers = S.fromList
+      . maybe [] (smallestSubsetBy preferred)
+      . LN.nonEmpty . S.toList
+      $ firstChoiceLosers
 
     preferred :: a -> a -> Ordering
-    preferred x0 x1 = compare (numMatch x0) (numMatch x1)
+    preferred x0 x1 = compare (numPrefers x0) (numPrefers x1)
       where
-        numMatch x = L.length . L.filter (== Just x) $ preferences
+        numPrefers x = length . LN.filter (== Just x) $ preferences
 
-        preferences :: [Maybe a]
-        preferences = map (ballotPreference . LN.toList) . LN.toList $ votes
-
-        ballotPreference :: [a] -> Maybe a
-        ballotPreference = L.find (`S.member` [x0, x1])
+        -- which candidate each voter prefers
+        preferences :: LN.NonEmpty (Maybe a)
+        preferences = LN.map (F.find (`S.member` [x0, x1])) $ votes
 
     voteCounts :: M.Map a Word
     voteCounts =
         M.unionsWith (+)
-        -- If there's a candidate with no first-place votes we need to include that
-      . (candidatesMinimumVoteCounts :)
+      . (candidatesStartingPoint :)
       . map (`M.singleton` 1)
       . LN.toList .
       LN.map LN.head $ votes
 
-    candidatesMinimumVoteCounts :: M.Map a Word
-    candidatesMinimumVoteCounts = (M.fromSet (const 0) (allVotedFor votes))
+    -- every candidate starts with zero votes
+    -- there are occasionally candidates with no first choice votes at all
+    -- we need to know about them
+    candidatesStartingPoint :: M.Map a Word
+    candidatesStartingPoint =
+      M.fromSet (const 0)
+      -- set of all candidates on anybody's ballot
+      . S.unions . LN.map (S.fromList . LN.toList) $ votes
 
 filterVotes :: (a -> Bool) -> LN.NonEmpty (LN.NonEmpty a) -> Maybe (LN.NonEmpty (LN.NonEmpty a))
 filterVotes f = catMaybes . LN.map (LN.nonEmpty . LN.filter f)
 
-allVotedFor :: Ord a => LN.NonEmpty (LN.NonEmpty a) -> S.Set a
-allVotedFor = S.unions . map (S.fromList . LN.toList) . LN.toList
-
 catMaybes :: LN.NonEmpty (Maybe a) -> Maybe (LN.NonEmpty a)
 catMaybes = LN.nonEmpty . M.catMaybes . LN.toList
 
-smallestSubset :: (t -> t -> Ordering) -> [t] -> [t]
-smallestSubset f = last . takeWhile (\xs -> (f (head xs) (last xs)) == O.EQ) . tail . L.inits . L.sortBy f
+-- smallest by the ordering function, not by size
+smallestSubsetBy :: forall t. (t -> t -> Ordering) -> LN.NonEmpty t -> [t]
+smallestSubsetBy f xs = LN.filter ((== EQ) . f (F.minimumBy f xs)) . LN.sortBy f $ xs