Haskell Quiz/Secret Santas/Solution Kuklewicz

{-# OPTIONS_GHC -fglasgow-exts -funbox-strict-fields #-}
-- Backtracking solution by Chris Kuklewicz <haskell@list.mightyreason.com>
-- http://haskell.org/haskellwiki/Haskell_Quiz/Secret_Santas
-- http://www.rubyquiz.com/quiz2.html
 
-- The chain of santas is cyclic.  This picks an initial person from
-- the remaining pool and creates a chain of secret santas until the
-- cycle is finished.  Then it is either done or it creates the next
-- cycle.  If a cycle cannot be completed then it backtracks and makes
-- a different choice of possible branches.
 
-- I believe this is a "fair" way of randomizing the assignments.
module Main where
import Control.Monad.Fix (fix)
import Data.Monoid
import Data.List (delete)
import System.Random (StdGen,newStdGen)
import MonadRandom(MonadRandom,Rand,getRandomR,evalRand) -- from http://haskell.org/haskellwiki/New_monads/MonadRandom
 
data Person = P {firstName,lastName,email :: !String}
  deriving (Show,Eq,Ord)
 
main = do
  input <- readFile "/tmp/input1"
  let people = map parsePerson (lines input)
  g <- newStdGen
  case assign people g of
    Nothing -> print "Failed"
    Just santas -> putStr (unlines (map showSanta santas))
 
parsePerson :: String -> Person
parsePerson line = P {firstName=a,lastName=b, email=c}
  where [a,b,c] = words line
 
showSanta (p1,p2) = (showPerson p1) ++ " -> " ++ (showPerson p2)
  where showPerson (P {firstName=a,lastName=b, email=c}) = unwords [a,b,c]
 
assign :: [Person] -> StdGen -> Maybe [(Person,Person)]
assign ps g = evalRand (assignAll ps) g
 
assignAll :: [Person] -> Rand StdGen (Maybe [(Person,Person)])
assignAll [] = return $ Just []
assignAll [_] = return Nothing
assignAll (initial:rest) = startCycle initial rest
 
startCycle :: Person -> [Person] -> Rand StdGen (Maybe [(Person,Person)])
startCycle _ [] = error "Cannot happen"
startCycle initial rest = do
  others <- shuffle (validOthers initial rest)
  let branches = map doBranch others
      doBranch latest = prepend initial latest $ continueCycle initial latest (delete latest rest)
  firstJust branches
 
continueCycle :: Person -> Person -> [Person] -> Rand StdGen (Maybe [(Person,Person)])
continueCycle initial previous rest | initial == previous = assignAll rest -- cycle was closed
                                    | otherwise = do
  others <- shuffle (validOthers previous (initial:rest))
  let branches = map doBranch others
      doBranch latest = prepend previous latest $ continueCycle initial latest (delete latest rest)
  firstJust branches
 
-- validOthers is used to exclude family members from the possibilities
validOthers (P {lastName=santaFamily}) = filter ((santaFamily/=).lastName)
 
-- helper function to build answer
prepend a b mv = do
  v <- mv 
  return $ case v of Nothing -> Nothing
                     Just cs ->  Just $ (a,b):cs
 
-- firstJust handles the backtracking on failure
firstJust [] = return Nothing
firstJust (mv:mvs) = do
  v <- mv
  case v of Nothing -> firstJust mvs
            Just _ -> return v
 
-- fairly randomize the order of a list
shuffle :: (MonadRandom m) => [a] -> m [a]
shuffle [] = return []
shuffle x@[_] = return x
shuffle xs = do
  let n = length xs
  rseq <- mapM (\i -> getRandomR (0,n-i)) [1..(n-1)]
  return (shuffle1 xs rseq)
 
-- Fair Shuffle copied from http://okmij.org/ftp/Haskell/perfect-shuffle.txt
 
-- A complete binary tree, of leaves and internal nodes.
-- Internal node: Node card l r
-- where card is the number of leaves under the node.
-- Invariant: card >=2. All internal tree nodes are always full.
data Tree a = Leaf !a | Node !Int !(Tree a) !(Tree a) deriving Show
 
build_tree = (fix grow_level) . (map Leaf)
	where
	     grow_level self [node] = node
	     grow_level self l = self $ inner l
 
	     inner [] = []
	     inner [e] = [e]
	     inner (e1:e2:rest) = (join e1 e2) : inner rest
 
	     join l@(Leaf _)       r@(Leaf _)       = Node 2 l r
	     join l@(Node ct _ _)  r@(Leaf _)       = Node (ct+1) l r
	     join l@(Leaf _)       r@(Node ct _ _)  = Node (ct+1) l r
	     join l@(Node ctl _ _) r@(Node ctr _ _) = Node (ctl+ctr) l r
 
-- given a sequence (e1,...en) to shuffle, and a sequence
-- (r1,...r[n-1]) of numbers such that r[i] is an independent sample
-- from a uniform random distribution [0..n-i], compute the
-- corresponding permutation of the input sequence.
 
shuffle1 elements rseq = shuffle1' (build_tree elements) rseq
	where
	     shuffle1' (Leaf e) [] = [e]
	     shuffle1' tree (r:r_others) = 
		let (b,rest) = extract_tree r tree
		in b:(shuffle1' rest r_others)
 
	     extract_tree 0 (Node _ (Leaf e) r) = (e,r)
	     extract_tree 1 (Node 2 (Leaf l) (Leaf r)) = (r,Leaf l)
	     extract_tree n (Node c (Leaf l) r) =
		let (e,new_r) = extract_tree (n-1) r
		in (e,Node (c-1) (Leaf l) new_r)
	     extract_tree n (Node n1 l (Leaf e)) 
			| n+1 == n1 = (e,l)
 
	     extract_tree n (Node c l@(Node cl _ _) r) 
			| n < cl = let (e,new_l) = extract_tree n l
				   in (e,Node (c-1) new_l r)
			| otherwise = let (e,new_r) = extract_tree (n-cl) r
				      in (e,Node (c-1) l new_r)