Heuristic stable roommates with k-person rooms by approval voting

Roommates.hs

import Data.Ord (Down(..), comparing)
import Data.List (permutations, sortBy)
import Data.Maybe (listToMaybe)
import qualified Data.Set as Set

-- | A matrix of approval votes. @forall (a :: 'Approvals'). a !! i !! j@
-- is the number of approval votes that person @i@ awarded to person @j@.
type Approvals = [[Int]]

-- | List of list of indices representing a partition of a list. Invariant:
--
-- > forall (x :: Partition). sort (concat x) == [0..pred (length (concat x))]
type Partition = [[Int]]

-- | Chunk size.
type Size = Int

-- | Approval score.
type Score = Int

-- | Example case with 2-person rooms:
--
-- Alice and Bob are friends; Alice is popular; Dan is unpopular.
--
--   * Alice approves of Bob and Charlie
--   * Bob approves of Alice and Dan
--   * Charlie approves of Alice and Bob
--   * Dan approves of Alice and Bob
--
-- The possible partitions and whether they’re stable are:
--
--   1. Alice and Bob (yes) / Charlie and Dan (no)
--   2. Alice and Dan (no) / Bob and Charlie (no)
--   3. Alice and Charlie (yes) / Bob and Dan (yes)
--
-- Of these, the one with the highest number of mutual approvals (2) is #3, so
-- the result of running @'solve' 2@ on this example is @[[0, 2], [1, 3]]@.
example :: Approvals
example =
  [ [1, 1, 1, 0]
  , [1, 1, 0, 1]
  , [1, 1, 1, 0]
  , [1, 1, 0, 1]
  ]

-- | Take the best solution, if there is one at all.
solve :: Size -> Approvals -> Maybe Partition
solve = fmap (fmap snd . listToMaybe . take 1) . solutions

-- | @'solutions' k a@ is the list of permutations of @a@ in descending order of
-- total number of mutual approval votes.
solutions :: Size -> Approvals -> [(Score, Partition)]
solutions k approvals = sortBy (comparing (Down . fst))
  [ (score, partition)
  | partition <- partitions k (length approvals)
  , let score = mutualApprovals approvals partition
  ]

-- | @'mutualApprovals' a p@ is the total number of approvals that any pair of
-- distinct roommates awarded to each other, in all rooms. Votes for yourself
-- are not counted.
mutualApprovals :: Approvals -> Partition -> Int
mutualApprovals approvals partition = sum
  [ x `min` y                    -- Could use sum, but min is # of shared votes
  | group <- partition           -- For each group (“room”) in partition
  , i <- group                   -- For each pair of people in the group
  , j <- group
  , i < j                        -- Only need to traverse one triangle
  , let x = approvals !! i !! j
  , x /= 0                       -- If i voted for j at all
  , let y = approvals !! j !! i
  , y /= 0                       -- And j voted for i at all
  ]

-- | @'partitions' k n@ generates (very inefficiently) a list of partitions into
-- chunks of length @k@ of the indices from 0 to @n - 1@.
partitions :: Size -> Int -> [Partition]
partitions k
  = unique
  . fmap (unique . fmap unique . chunksOf k)
  . permutations
  . enumFromTo 0
  . pred

-- | @'chunksOf' n xs@ partitions @xs@ into chunks of length @n@. Fails if the
-- input list length is not evenly divisible by the chunk length.
chunksOf :: Int -> [a] -> [[a]]
chunksOf n = go []
  where
    go acc xs = case splitAt n xs of
      ([], []) -> reverse acc
      (as, bs)
        | length as == n -> go (as : acc) bs
        | otherwise -> error "chunksOf: length not divisible by chunk length"

-- | Slightly more efficient than 'Data.List.nub'.
unique :: Ord a => [a] -> [a]
unique = Set.toList . Set.fromList