72 lines
2.5 KiB
Haskell
72 lines
2.5 KiB
Haskell
module InstantRunoff where
|
|
import qualified Data.List as L
|
|
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.Foldable as F
|
|
|
|
solve :: forall a. Ord a => [a] -> LN.NonEmpty (LN.NonEmpty a) -> [S.Set a]
|
|
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)
|
|
-- all candidates on any submited ballot
|
|
. S.unions . LN.map (S.fromList . LN.toList) $ votes
|
|
|
|
|
|
rank :: forall a. Ord a => LN.NonEmpty (LN.NonEmpty a) -> [S.Set a]
|
|
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
|
|
. maybe [] (smallestSubsetBy preferred)
|
|
. LN.nonEmpty . S.toList
|
|
$ firstChoiceLosers
|
|
|
|
preferred :: a -> a -> Ordering
|
|
preferred x0 x1 = compare (numPrefers x0) (numPrefers x1)
|
|
where
|
|
numPrefers x = length . LN.filter (== Just x) $ preferences
|
|
|
|
-- 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 (+)
|
|
. (M.fromSet (const 0) candidates :)
|
|
. map (`M.singleton` 1)
|
|
. LN.toList .
|
|
LN.map LN.head $ votes
|
|
|
|
-- set of all candidates on anybody's ballot
|
|
candidates :: S.Set a
|
|
candidates = 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 = InstantRunoff.catMaybes . LN.map (LN.nonEmpty . LN.filter f)
|
|
|
|
catMaybes :: LN.NonEmpty (Maybe a) -> Maybe (LN.NonEmpty a)
|
|
catMaybes = LN.nonEmpty . M.catMaybes . LN.toList
|
|
|
|
-- 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
|