Gale-Shapely in Haskell

matching.hs

{-# LANGUAGE FlexibleContexts #-}

--- Some experiments vis-a-vis the Gale-Shapely algorithm.

import Data.List
import Control.Monad
import Data.Array
import Data.Array.MArray
import Data.Array.ST
import Data.Ord
import Data.List
import System.Random
import System.Random.Shuffle
import Data.Char
import Control.Monad.Random.Class
import Debug.Trace
import Text.Printf
import qualified Data.Map.Strict as Map

import Control.Monad.ST

--- This matching algorithm does not work.
---

eqFst (a,_) (b,_) = a == b

grouped :: Ord b => [ (b,a) ] -> [ (b, [a]) ]
grouped pairs = [ (b, as) | g <- groupBy eqFst $ sortBy (comparing fst) pairs,
                            let b = fst (head g)
                                as = map snd g ]
match :: (Ix a, Ix b, Show a, Show b)
      => Array (a,b) Int
      -> Array (b,a) Int
      -> [a]
      -> [b]
      -> [ (a,b) ]
match aPrefs bPrefs as bs =
  let groups = grouped [ (b,a) | a <- as, let b = aChoice a ]
       where tr x = trace ("groups: " ++ show x) x

      pairs = [ (a,b) | (b, as) <- groups, let a = bChoice b as ]
  in pairs
  where
    bAvailable b = elem b bs
    aChoice a    = minimumBy (comparing (\b -> aPrefs ! (a,b))) bs
      where tr x = trace msg x
              where msg = "aChoice " ++ show a ++ " = " ++ show x
    bChoice b as = minimumBy (comparing (\a -> bPrefs ! (b,a))) as

matchAll
  :: (Ix a, Ix b, Show a, Show b)
  => Array (a,b) Int
  -> Array (b,a) Int
  -> [a]
  -> [b]
  -> [ [ (a,b) ] ]
matchAll aPrefs bPrefs as bs = iterate go []
  where
    go paired = paired ++ newpairs
      where newpairs = match aPrefs bPrefs as' bs'
            as' = as \\ (map fst paired)
            bs' = bs \\ (map snd paired)

---
---

generatePrefs :: MonadRandom m => Int -> m (Array (Char,Char) Int, Array (Char,Char) Int, String, String)
generatePrefs n = do
  let as = take n ['A'..]
      aend = last as
      bs = take n ['a'..] 
      bend = last bs
  aprefs <- replicateM n (shuffleM [1..n])
  bprefs <- replicateM n (shuffleM [1..n])
  let aPrefs = listArray (('A','a'),(aend,bend)) $ concat aprefs
      bPrefs = listArray (('a','A'),(bend,aend)) $ concat bprefs
  return (aPrefs, bPrefs, as, bs)

test :: (Ix a, Ix b) => a -> b -> Array (a,b) Int -> Int
test a b arr = arr ! (a,b)

test2 :: (Ix a, Ix b) => a -> [b] -> Array (a,b) Int -> b
test2 a bs arr = minimumBy (comparing (\b -> arr ! (a,b))) bs

eShow x = [x]

inversePerm :: Int -> [Int] -> [Int]
inversePerm n ranks =
  let arr = array (1,n) [ (r,k) | (k,r) <- zip [1..] ranks ]
  in elems arr

showPrefs arr = do
  let ((lo,l),(hi,h)) = bounds arr
      n = fromEnum hi - fromEnum lo + 1
      toB k = [ toEnum ( fromEnum l + k - 1 ) ]
  forM_ [lo..hi] $ \a -> do
    let as = eShow a
        row = intercalate " " $ map toB $ inversePerm n $ [ arr ! (a,b) | b <- [l..h] ]
    putStrLn $ as ++ " - " ++ row

test3 = do
  g <- getStdGen
  (aPrefs, bPrefs, _, _) <- generatePrefs 10
  showPrefs aPrefs
  showPrefs bPrefs

test4 = test5 4

showPairs pairs = intercalate ", " [ [a] ++ "-" ++ [b] | (a,b) <- pairs ]

takeUpTo p [] = []
takeUpTo p (x:xs) | p x = [x]
                  | otherwise = x : takeUpTo p xs

test5 n = do
  g <- getStdGen
  (aPrefs, bPrefs, as, bs) <- generatePrefs n
  showPrefs aPrefs
  showPrefs bPrefs
  mapM_ (putStrLn.showPairs) $ takeUpTo (\z -> length z >= n) $ matchAll aPrefs bPrefs as bs

allpairs xs = [ (p,q) | (p:ps) <- tails xs, q <- ps ]

isStable aPrefs bPrefs pairs =
  filter (uncurry unstable) (allpairs pairs)
    where unstable (a1,b1) (a2,b2) = (aprefers a1 b2 b1 && bprefers b2 a1 a2)
                                     ||
                                     (aprefers a2 b1 b2 && bprefers b1 a2 a1)
          aprefers a b1 b2 = aPrefs ! (a,b1) < aPrefs ! (a,b2)
          bprefers b a1 a2 = bPrefs ! (b,a1) < bPrefs ! (b,a2)

checkStable aPrefs bPrefs ((a1,b1),(a2,b2))
  = (aprefers a1 b2 b1, bprefers b2 a1 a2, aprefers a2 b1 b2, bprefers b1 a2 a1)
  where
    aprefers a b1 b2 = aPrefs ! (a,b1) < aPrefs ! (a,b2)
    bprefers b a1 a2 = bPrefs ! (b,a1) < bPrefs ! (b,a2)

explain aPrefs bPrefs ((a1,b1), (a2,b2)) = do
  putStrLn $ printf "pairing %c-%c v. %c-%c" a1 b1 a2 b2
  let a1b1 = aPrefs ! (a1,b1)
      a1b2 = aPrefs ! (a1,b2)
      a2b1 = aPrefs ! (a2,b1)
      a2b2 = aPrefs ! (a2,b2)
  let b1a1 = bPrefs ! (b1,a1)
      b1a2 = bPrefs ! (b1,a2)
      b2a1 = bPrefs ! (b2,a1)
      b2a2 = bPrefs ! (b2,a2)
  go a1 b2 b1 a1b2 a1b1
  go b2 a1 a2 b2a1 b2a2

  go a2 b1 b2 a2b1 a2b2
  go b1 a2 a1 b1a2 b1a1

  where
    go x y1 y2 r1 r2 =
      putStrLn $ printf "%c prefers %c to %c = %d < %d = %s" x y1 y2 r1 r2 (show (r1 < r2))

test6 n = do
  g <- getStdGen
  (aPrefs, bPrefs, as, bs) <- generatePrefs n
  showPrefs aPrefs
  showPrefs bPrefs
  let matchings = takeUpTo (\z -> length z >= n) $ matchAll aPrefs bPrefs as bs
      matching = last matchings
  mapM_ (putStrLn.showPairs) matchings
  forM_ (isStable aPrefs bPrefs matching) $ \(m1,m2) -> do
    explain aPrefs bPrefs (m1,m2)
--  mapM_ print $ map (checkStable aPrefs bPrefs) (allpairs matching)

--- This algorithm does work.
---

galeShapely :: Array (Char,Char) Int -> Array (Char,Char) Int -> [Char] -> [Char] -> [ (Char,Char) ]
galeShapely aPrefs bPrefs as bs = runST $ do
  let n = length as
      aname = listArray (1,n) as
      bname = listArray (1,n) bs

      bTojMap = Map.fromList (zip bs [(1::Int)..])
      bToj b = Map.findWithDefault 0 b bTojMap

  proposals <- newListArray (1,n) $ [ inversePerm n $ [ aPrefs ! (a,b) | b <- bs ]
                                          | a <- as
                                    ] :: ST s (STArray s Int [Int])

  apartner  <- newListArray (1,n) (repeat 0)  :: ST s (STArray s Int Int)
  bpartner  <- newListArray (1,n) (repeat 0)  :: ST s (STArray s Int Int)

  let advance k = go [ 1 + (mod (k+i) n) | i <- [1..n] ]
        where go []     = return Nothing
              go (0:_)  = error "k = 0"
              go (k:ks) = do v <- readArray apartner k
                             if v == 0
                                then do props <- readArray proposals k
                                        case props of
                                          (j:js) -> do writeArray proposals k js
                                                       return $ Just (k, j)
                                          _      -> go ks
                                else go ks
  let loop k = do
        m <- advance k
        case m of
          Nothing      -> return ()
          Just (k1, j) -> do              -- k1 should propose to j
            k0 <- readArray bpartner j
            if k0 == 0
              then do writeArray apartner k1 j
                      writeArray bpartner j  k1
              else let r0 = bPrefs ! (bname!j, aname!k0)
                       r1 = bPrefs ! (bname!j, aname!k1)
                   in
                   if r1 < r0
                     then do writeArray apartner k1 j
                             writeArray bpartner j  k1
                             writeArray apartner k0 0
                     else return ()
            loop k1

  loop 1
  js <- getElems apartner
  ps <- getElems proposals
  let tr x = trace msg x
         where msg = unlines $ [ "js = " ++ show x] ++ ( map show ps )
  return $ [ ( a, bname!j ) | (a, j) <- zip as js ]

testg1 n = do
  g <- getStdGen
  (aPrefs, bPrefs, as, bs) <- generatePrefs n
  showPrefs aPrefs
  showPrefs bPrefs
  let matching = galeShapely aPrefs bPrefs as bs
  (putStrLn.showPairs) matching
  let unstables = isStable aPrefs bPrefs matching
  case unstables of
    [] -> putStrLn "No unstable pairs."
    _  -> forM_ unstables $ explain aPrefs bPrefs