Programming Praxis – Who Owns The Zebra?

In Today’s Programming Praxis problem we have to solve a logic puzzle. The provided solution uses a 182-line logic programming library and then takes 36 lines to solve the problem. I didn’t feel like porting 182 lines from Scheme to Haskell, so I rolled my own solution.  It’s going to be a slightly longer one than usual though, so let’s dive right in.

Our imports:

import Data.List
import qualified Data.Map as M

We’re going to handle the constraints by applying them to a two-dimensional grid. One axis holds the position of the house (first, second, etc.) and the other the various properties (nationality, color, etc.). Each cell holds the remaining options for that combination of house and property. By applying constraints we’re going to remove options until each cell has only one option left. It’s a bit like sudoku puzzles if you think about it.

type Grid = M.Map String (M.Map Int [String])

In the problem we have four types of constraints, which we encode in an ADT:

data Constraint = Link (String, String) (String, String)
                | PosLink (String, String) Int
                | NextTo (String, String) (String, String)
                | RightOf (String, String) (String, String)
                deriving Eq

A convenience type to keep the type signatures a bit easier to read:

type Solver = ([Constraint], Grid)

Adding a constraint to a solver is trivial:

addConstraint :: Constraint -> Solver -> Solver
addConstraint c (cs, g) = (c : cs, g)

This function abstracts out some common logic. It removes options from the grid if the conditions to do so have been met.

removeIf :: (String, String) -> (String, String) ->
    [String -> String -> Int -> Grid -> Bool] -> Grid -> Grid
removeIf (f1, v1) (f2, v2) cs g = M.adjust (M.mapWithKey (\k ->
    if and [c f1 v1 k g | c <- cs] then delete v2 else id)) f2 g

Like removeIf, notAt abstract out some common code. It checks if a given value is still an option for the given property in another house.

notAt :: (Int -> Int) -> String -> String -> Int -> Grid -> Bool
notAt f f1 v1 i g = M.notMember (f i) (g M.! f1) ||
                    notElem v1 (g M.! f1 M.! (f i))

With that out of the way, the function to apply a constraint looks like this. Since most constraints work in two directions, we have to apply them in both directions.

runConstraint :: Constraint -> Grid -> Grid
runConstraint (Link a b) = removeIf a b conds . removeIf b a conds
    where conds = [(\f1 v1 k -> notElem v1 . (M.! k) . (M.! f1))]
runConstraint (PosLink (f1,v1) i) =
    M.adjust (M.update (const $ Just [v1]) i) f1
runConstraint (NextTo a b)  = removeIf a b [notAt pred, notAt succ]
runConstraint (RightOf a b) = removeIf a b [notAt pred] .
                              removeIf b a [notAt succ]

adjustOthers applies a function to all elements of a map except the given one, which we need for the next function.

adjustOthers :: Eq k => (v -> v) -> k -> M.Map k v -> M.Map k v
adjustOthers f k = M.mapWithKey (\k' v -> if k' == k then v else f v)

If a house has only one option left for a property than we can remove that option from all the other houses. Similarly, if a house is the only one that still has a certain option, we can remove the other options for that property.

simplify :: Grid -> Grid
simplify g = foldr ($) (M.mapWithKey (\_ v ->
    M.mapWithKey (\i x -> let d = x \\ concat (M.elems $ M.delete i v)
                          in if length d == 1 then d else x) v) g)
    [ M.adjust (adjustOthers (\\ take 1 x) i) f
    | (f, v) <- M.assocs g, (i, x) <- M.assocs v, length x == 1]

run simply runs all the constraints once.

run :: Solver -> Solver
run (cs, g) = (cs, simplify $ foldr runConstraint g cs)

Once all the constraints have been run, we might have fewer options available than we did at the beginning, which might open up new possibilities for more removal. apply keeps applying all the constraints until no further progress is made.

apply :: Solver -> Solver
apply = head . head . dropWhile (null . tail) . group . iterate run

If we had enough constraints to solve the problem with just constraint propagation we could stop here. Unfortunately, this doesn’t work on the problem we have to solve.  While it significantly reduces the available options, it can’t give a complete solution. So we’re going to have to do what any self-respecting logician would do in such a scenario: guess. If a property still has multiple options we choose one of them and see if we can solve it then. If not, we try the next option, or we do the same thing for the next property if none of the guesses helps solve the problem.

If any property still has more than one option the problem is not solved.

solved :: M.Map k (M.Map k' [v]) -> Bool
solved g = and [False | (_, v)  <- M.assocs g,
                        (_, xs) <- M.assocs v, length xs /= 1]

solve takes care of the guesswork, and also reformats the output to be more readable.

solve :: Solver -> [String]
solve s = map (unwords . map head) . transpose . map (M.elems) .
          M.elems $ head [ r | let (cs, g) = apply s,
          (f, v) <- M.assocs $ g, (i, xs) <- M.assocs v, x <- xs,
          let (_, r) = apply (cs, M.adjust (M.adjust (const [x]) i) f g),
          solved r ]

And there we have our constraint solver. Now for the problem. First we create the grid with all the options:

grid :: Grid
grid = M.fromList . zip (words "owner brand drink pet color") $
    map (M.fromList . zip [1..] . replicate 5)
    [words "Englishman Ukranian Norwegian Japanese Spaniard",
     words "Old_Gold Kools Chesterfields Lucky_Strike Parliaments",
     words "Coffee Tea Milk Orange_Juice Water",
     words "Dog Snails Horse Fox Zebra",
     words "Red Green Ivory Yellow Blue"]

Next we add all our constraints.

problem :: Solver
problem = foldr addConstraint ([], grid)
    [Link    ("owner", "Englishman")    ("color", "Red"),
     Link    ("owner", "Spaniard")      ("pet",   "Dog"),
     Link    ("drink", "Coffee")        ("color", "Green"),
     Link    ("owner", "Ukranian")      ("drink", "Tea"),
     RightOf ("color", "Ivory")         ("color", "Green"),
     Link    ("brand", "Old_Gold")      ("pet",   "Snails"),
     Link    ("brand", "Kools")         ("color", "Yellow"),
     PosLink ("drink", "Milk")          3,
     PosLink ("owner", "Norwegian")     1,
     NextTo  ("brand", "Chesterfields") ("pet",   "Fox"),
     NextTo  ("brand", "Kools")         ("pet",   "Horse"),
     Link    ("brand", "Lucky_Strike")  ("drink", "Orange_Juice"),
     Link    ("owner", "Japanese")      ("brand", "Parliaments"),
     NextTo  ("owner", "Norwegian")     ("color", "Blue")]

And finally we print the solution.

main :: IO ()
main = mapM_ putStrLn $ solve problem

That brings the total to 36 lines for the solver and 23 for the problem, and it runs in about 60 ms. I’d say that will do nicely.

Tags: constraint, Haskell, kata, logic, praxis, programming, zebra