Solver for puzzles in 999

Solver.hs

{-# LANGUAGE ScopedTypeVariables, ViewPatterns, TupleSections #-}

module Solver999 where

import Data.List
import Control.Monad
import Data.Set (Set)
import qualified Data.Set as S
import Data.Map (Map)
import qualified Data.Map as M
import Data.Vector (Vector)
import qualified Data.Vector as V

--- Splits a number into a list of decimal digits.
decimalDigits :: Int -> [Int]
decimalDigits n
  -- No upper digits left, return the current one.
  | left == 0 = [cur]
  -- Get digits of right part of the number and append this one.
  | otherwise = decimalDigits left ++ [cur]
  -- Get integer division and modulo of a number
  -- left: integer division, cur: modulo
  where (left, cur) = n `divMod` 10

-- Get a digital root for given digits.
digitalRoot :: [Int] -> Int
digitalRoot digits
  | s < 10 = s
  | otherwise = digitalRoot $ decimalDigits s
  where s = sum digits

-- Check validity of a group w.r.t. door.
acceptableGroup :: Int -> Set Int -> Bool
acceptableGroup d xs = len >= 3 && len <= 5 && digitalRoot (S.toList xs) == d
  where len = S.size xs

-- Update an item in vector by index, using given function.
updateVector :: Int -> (a -> a) -> Vector a -> Vector a
updateVector i upd v = v V.// [(i, upd (v V.! i))]

-- Get all possible splits of unique items to N baskets.
-- For example, possible splits of [1, 2] into baskets ["a", "b"] are:
--  * [("a", [1, 2]), ("b", [])]
--  * [("a", [1]),    ("b", [2])]
--  * [("a", [2]),    ("b", [1])]
--  * [("a", []),     ("b", [1, 2])]
-- Returns maps from basket to items in it.
splitToBaskets :: (Ord k, Ord a) => Set k -> Set a -> [Map k (Set a)]
splitToBaskets basketNames = make . S.toList
  -- If there are no items, return empty baskets.
  where make [] = [M.fromSet (const S.empty) basketNames]
        -- If there are items left, get all possible splits without one item and add this item to each basket in turn.
        make (x:xs) = [ M.adjust (S.insert x) k baskets -- Return baskets with x inserted into basket with name k...
                      | baskets <- make xs -- For all splits without x...
                      , k <- S.toList basketNames -- and for all basket names k
                      ]

-- Get all possible splits of unique items to arbitrary sized groups.
-- For example, possible groups for [1, 2, 3] are:
--  * [[1], [2], [3]]
--  * [[1, 2], [3]]
--  * [[2], [1, 3]]
--  * [[3], [1, 2]]
--  * [[1, 2, 3]]
splitToGroups :: forall a. Ord a => Set a -> [[Set a]]
splitToGroups = map V.toList . make . S.toList
  where make :: [a] -> [Vector (Set a)]
        make [] = [V.singleton S.empty]
        make [x] = [V.singleton (S.singleton x)]
        make (x:xs) = concatMap makeOne $ make xs
          where makeOne groups = V.cons (S.singleton x) groups : map insertOne [0..V.length groups - 1]
                  where insertOne i = updateVector i (S.insert x) groups

-- Given a set of people and set of doors, determine possible all-people solutions.
assignDoors :: Set Int -> Set Int -> [Map Int [Set Int]]
assignDoors doors people = concatMap (M.traverseWithKey tryGroups) $ splitToBaskets doors people
  where tryGroups door = filter (all (\xs -> S.null xs || acceptableGroup door xs)) . splitToGroups

-- Combinations of N objects by K.
combinations :: Ord a => Int -> Set a -> [Set a]
combinations n0 = make n0 . S.toList
  where make 0 _ = [S.empty]
        make k lst = [ S.insert x rest
                     | x:xs <- tails lst
                     , rest <- make (k - 1) xs
                     ]

-- Breadth-first search for a pins puzzle solution.
pinsSolution :: [(Set Int, Set Int)]
-- At each step we maintain a map of possible pin states (which ones are activated) to solution path.
-- We start with a single entry: all pins disable, with empty solution path (you don't need to do anything to get into that state).
-- Each step we derive a map of new states and merge it with old ones. Notice that this way if two solution paths converges at some point into one state they will be joined.
pinsSolution = refine $ M.singleton (V.snoc (V.replicate 8 False) True) []
  where refine states = case M.lookup (V.replicate 9 True) states of
          -- We check if there is a path that gets us to all pins activated. If there is, we stop the process and return the path.
          -- Reversed because we add steps to the beginning of the path list during procedure.
          Just v -> reverse v
          _ -> refine $ M.fromList [ (update p1 $ update p2 state, step : way)
                                  -- Get old states
                                  | (state, way) <- M.toList states
                                  -- Try all possible steps
                                  , step@(p1, p2) <- combs
                                  ]

        -- This function updates the state given old one and a performed step.
        update p = updateVector (digitalRoot (S.toList p) - 1) not

        pins = S.fromList [1..9]
        combs = [ (c1, c2)
                | c1 <- combinations 3 pins
                , c2 <- combinations 3 (pins S.\\ c1)
                ]

-- Solution for the magic square puzzle.
pins2Solution :: [Vector Int]
pins2Solution = filter check $ map (\(splitAt 4 -> (a1, a2)) -> V.fromList (a1 ++ [5] ++ a2)) $ permutations ([1..4] ++ [6..9])
  where check vec = all ((== 15) . sum . map (vec V.!)) lineIndices
        lineIndices = map (map (\(x, y) -> y * 3 + x)) lines
        lines =    [ map (x, ) [0..2]
                   | x <- [0..2]
                   ]
                ++ [ map (, y) [0..2]
                   | y <- [0..2]
                   ]
                ++ [ [(0,0), (1,1), (2,2)]
                   , [(2,0), (1,1), (0,2)]
                   ]

studySolution :: [(Int, Int)]
studySolution = refine $ M.singleton initialState []
  where refine states = case find (checkBoard . fst) $ M.toList states of
          Just (_, v) -> reverse v
          _ -> refine $ M.fromList [ (update step state, step : way)
                                  | (state, way) <- M.toList states
                                  , step@(x, y) <- combs
                                  ]

        size = 3
        board = V.fromList [ 13, 6, 4
                           , 1, 15, 5
                           , 9, 11, 12
                           ]
        initialState = V.fromList [ True, True, False
                                  , True, True, True
                                  , True, True, False
                                  ]
        wanted = [ ([0, 1, 2], 10)
                 , ([3, 4, 5], 21)
                 , ([6, 7, 8], 21)
                 , ([0, 3, 6], 10)
                 , ([1, 4, 7], 21)
                 , ([2, 5, 8], 21)
                 ]

        checkBoard state = all (\(path, s) -> sum (map (\i -> if state V.! i then board V.! i else 0) path) == s) wanted

        update step state = state V.// map swapPoint (neighbours step)
          where swapPoint p@(x, y) =
                  let i = y * size + x
                  in (i, not $ state V.! i)
        neighbours (x, y) = filter (\(x, y) -> x >= 0 && x < size && y >= 0 && y < size)
                            [ (x - 1, y), (x, y), (x + 1, y)
                            , (x, y - 1),         (x, y + 1)
                            ]

        combs = [ (x, y) | x <- [0..size - 1], y <- [0..size - 1] ]