StableMarriage.hs

{-#LANGUAGE TemplateHaskell, NoImplicitPrelude, StandaloneDeriving #-}

module StableMarriage
  ( -- * Usage
    -- $usage
    Id
  , Rank
  , stableMarriage
    -- * Internals
    -- ** Partner
  , Partner
    -- *** Lenses
  , id
  , preferences
  , proposals
  , partner
    -- ** Gender
    -- $gender
  , Identity
  , Male
  , Female
  , Man(..)
  , Woman(..)
    -- ** Algorithm
  , accepts
  , marry
  , marryMan
  , marryMen
  ) where

import Control.Comonad.Identity
import Control.Lens
import Control.Monad.State.Lazy
import Data.Array.Lens
import Data.Array.Unboxed
import Data.Maybe
import Prelude hiding (id)

-- $usage This module provides an implementation of the Gale-Shapley
-- Algorithm for solving the Stable Marriage Problem. For more information, see:
-- <http://en.wikipedia.org/wiki/Stable_marriage_problem>.
--
-- In the Stable Matching Problem, every 'Man' and every 'Woman' is identified 
-- by an 'Id' unique to their group and maintains a 'Rank'ing for each member of
-- the opposite group. This module exports the function 'stableMarriage', which
-- accepts that information encoded as two lists of pairs of 'Id' and @['Id']@.
--
-- For example:
--
-- > men   = [ (1, [2, 1])
-- >         , (2, [1, 2]) ]
-- > women = [ (1, [1, 2])
-- >         , (2, [2, 1]) ]
-- > stable = stableMarriage men women

type Id   = Int
type Rank = Int

-- $gender 'Male' and 'Female' are gendered 'Identity' 'Functor's for annotating
-- 'Partner's. The 'Monad' and 'Comonad' instances of 'Identity' allow use to
-- 'return' and 'extract' gendered values (for example, 'Id's).
--
-- 'Man' and 'Woman' are shorthand for @'Partner' 'Male' 'Female'@ and
-- @'Partner' 'Female' 'Male'@, respectively.

type Male   = Identity
type Female = Identity

deriving instance Eq  a => Eq  (Identity a)
deriving instance Ord a => Ord (Identity a)
deriving instance Ix  a => Ix  (Identity a)

-- | Every @'Partner' f g@ is uniquely identified by some @f 'Id'@ and
-- maintains a 'Rank' and proposal order for all other @'Partner' g f@s.
data Partner f g = Partner
  { _id          :: f Id                   -- ^ Id, unique among for all @f 'Id'@
  , _preferences :: UArray (g Id) Rank     -- ^ Rankings of all @'Partner' g f@s
  , _proposals   :: [g Id]                 -- ^ Proposal order of all @'Partner' g f@s
  , _partner     :: Maybe (Partner g f) }  -- ^ Married @'Partner' g f@, if any

makeLenses ''Partner

newtype Man   = Man   (Partner Male Female)
newtype Woman = Woman (Partner Female Male)

-- | A 'Woman' accepts a 'Man' if she has no partner or if she ranks him higher
-- than her current partner.
accepts :: Woman -> Man -> Bool
(Woman w) `accepts` (Man m) =
  let m'    = w ^. partner
      ps    = w ^. preferences
      rank  = ps ! (m  ^. id)
      rank' = ps ! (fromJust m' ^. id)
  in  isNothing m' || rank < rank'

-- | Marries a 'Woman' and 'Man' and divorces the previous partner, if any.
marry :: Woman -> Man -> (Woman, Maybe Man)
(Woman w) `marry` (Man m) =
  let (r, w') = partner <<.~ Just m $ w
      r'      = over _just (proposals %~ tail) r
  in  (Woman w', fmap Man r')

-- | /O(n)/. Marries a 'Man' to a 'Woman' and returns the 'Man' he replaced, if
-- any. Maintains the state of each 'Woman'.
marryMan :: Man -> State (Array (Female Id) Woman) (Maybe Man)
marryMan (Man m) = do
    women <- get
    let (i:_, m') = proposals <%~ dropRejectors women (Man m) $ m
        (w', r)        = (women ! i) `marry` Man m'
    put $ ix i .~ w' $ women
    return r
  where
    dropRejectors women m =
      dropWhile $ \i ->
        let w = women ! i
        in  not $ w `accepts` m

-- | /O(n²)/. Marries each 'Man' in a list to a 'Woman', recursing if necessary,
-- and returns each married 'Woman'. Maintains the state of each 'Woman'.
marryMen :: [Man] -> State (Array (Female Id) Woman) [Woman]
marryMen men = do
  rejects <- liftM catMaybes $ forM men marryMan
  state <- get
  if null rejects
    then return $ elems state
    else marryMen rejects

-- | /O(n²)/. Run the Stable Marriage algorithm.
stableMarriage :: [(Id, [Id])] -> [(Id, [Id])] -> [(Id, Id)]
stableMarriage ms ws =
    let men          = fmap (Man . fromTuple) ms
        bounds       = ( return . minimum $ map fst ws
                       , return . maximum $ map fst ws )
        women        = array bounds
                     $ map (\w -> (return $ fst w, Woman $ fromTuple w)) ws
        matchedWomen = evalState (marryMen men) women
    in  map toTuple matchedWomen
  where
    fromTuple (id, preferences) =
      let preferences' = map return preferences
      in Partner
         { _id          = return id
         , _preferences = array (return 1, return . fromIntegral $ length preferences)
                        $ zip preferences' [1..]
         , _proposals   = preferences'
         , _partner     = Nothing }
    toTuple (Woman w) = ( extract $ fromJust (w ^. partner) ^. id
                        , extract $ w ^. id )