【Haskell】SATソルバを使ってスリザーリンクを解く

-- This code requires SAT solver module,
-- e.g. https://gist.github.com/gatlin/1755736

import Data.Char (isDigit, ord)
import Data.List (genericLength, group, groupBy, sort)
import Data.Map (alter, empty, (!))
import Sat (solve)

problem = ["-0-1--1-",
           "-3--23-2",
           "--0----0",
           "-3--0---",
           "---3--0-",
           "1----3--",
           "3-13--3-",
           "-0--3-3-"]

main = do
  case solveSlitherLink problem of
    Nothing -> putStrLn "Impossible"
    Just answer -> putStrLn $ showAnswer problem answer

solveSlitherLink problem = solveUntil (oneLoop w) $ makeSlitherLinkConstraints problem
  where w = genericLength $ head problem

makeSlitherLinkConstraints problem = concat $ gridCnstss ++ numCnstss
  where gridCnstss = gridConstraints w h
        numCnstss = numberConstraints problem
        w = genericLength $ head problem
        h = genericLength problem

-- Whether the given answer makes one loop.
oneLoop :: Integer -> [Integer] -> Bool
oneLoop w answer = length edges == length (search $ fst $ head edges)
  where edges = map (toGridPoints) $ filter (> 0) answer
        toGridPoints var | isHorzEdge var = horz var
                         | otherwise      = vert var
        isHorzEdge v = (v - 1) `mod` lw < w
        horz v = (((v - 1) `mod` lw, (v - 1) `div` lw),
                  ((v - 1) `mod` lw + 1, (v - 1) `div` lw))
        vert v = (((v - 1) `mod` lw - w, (v - 1) `div` lw),
                  ((v - 1) `mod` lw - w, (v - 1) `div` lw + 1))
        lw = 2 * w + 1
        graph = foldl update empty edges
        update m (c1, c2) = alter (f c2) c1 $ alter (f c1) c2 m
        f c Nothing = Just [c]
        f c (Just o) = Just (c:o)
        search pt = loop [pt] $ head (graph ! pt)
        loop path p | p `elem` path = path
                    | otherwise     = loop (p:path) $ next p (head path)
        next p o = head $ filter (/= o) $ (graph ! p)

---- Cell constraints.
-- 4 Cell edges must be the number of the cell.

numberConstraints :: [String] -> [[[Integer]]]
numberConstraints problem = map (cellConstraints w) $ map (\(n,pt) -> (ord n - ord '0', pt)) $ cellNums
  where cellNums = filter (isDigit . fst) $ zip (concat problem) $ allCells
        allCells = [(x,y) | y <- [0..h-1], x <- [0..w-1]]
        w = genericLength $ head problem
        h = genericLength problem

cellConstraints w (n,(x,y)) = cnst
  where cnst = dnf2cnf edges $ cellEdgePatterns !! n
        edges = [u, u + w, u + w + 1, u + 2 * w + 1]
        u = y * (2 * w + 1) + x + 1

cellEdgePatterns :: [[[Integer]]]
cellEdgePatterns = map (\n -> filter (edgeNum n) ps) [0..4]
  where edgeNum n = (== n) . length . filter (> 0)
        ps = multiply $ replicate 4 [-1, 1]

---- Grid constraints.
-- Every grid point have 0 or 2 edges.

gridConstraints :: Integer -> Integer -> [[[Integer]]]
gridConstraints w h = map mulCoeffs [gridEdges w h x y | (x,y) <- allGridPoint]
  where mulCoeffs ls = dnf2cnf ls $ gridPointEdgePatterns !! (length ls)
        allGridPoint = [(x, y) | y <- [0..h], x <- [0..w]]

gridEdges :: Integer -> Integer -> Integer -> Integer -> [Integer]
gridEdges w h x y = concat [u, l, r, d]
  where u | y == 0    = []
          | otherwise = [p - w - 1]
        l | x == 0    = []
          | otherwise = [p - 1]
        r | x == w    = []
          | otherwise = [p]
        d | y == h    = []
          | otherwise = [p + w]
        p = y * (2 * w + 1) + x + 1

gridPointEdgePatterns :: [[[Integer]]]
gridPointEdgePatterns = map convert [0..4]
  where convert = filter has0or2edges . multiply . flip replicate [-1, 1]
        has0or2edges = is0or2 . length . filter (> 0)
        is0or2 n = n == 0 || n == 2

-- Convert SAT result to visible form.
showAnswer :: [String] -> [Integer] -> String
showAnswer problem answer = unlines $ map concat $ sandwich horz cells
  where frame = map (splitAt w) $ pack (2 * w + 1) $ sortOf abs answer
        w = length $ head problem
        cells = zipWith sandwich (map (map toVert. snd) frame) $ map label pproblem
        label = map (\c -> [' ', c, ' '])
        horz  = map (sandwich (repeat "+") . map toHorz . fst) frame
        pproblem = map (map (\c -> if c `elem` "0123" then c else ' ')) problem
        toHorz x | x > 0     = "---"
                 | otherwise = "   "
        toVert x | x > 0     = "|"
                 | otherwise = " "

solveUntil :: ([Integer] -> Bool) -> [[Integer]] -> Maybe [Integer]
solveUntil f constraints =
  case solve constraints of
    Nothing -> Nothing
    Just answer -> if f answer
                     then Just answer
                     else solveUntil f $ exclude answer : constraints
  where exclude answer = map negate answer

-- DNF (Disjunctive Normal Form, ex. (A and B) or (C and D) or ...) to
-- CNF (Conjunctive Normal Form, ex. (A or C) and (A or D) and ...)
--dnf2cnf :: Num b => [b] -> [[Integer]] -> [[Integer]]
dnf2cnf vars dnf = map (map ref) reducedCnf
  where reducedCnf = restoreVar $ removeTrueClauses $ map groupVars cnf
        cnf = multiply $ map (zipWith (*) [1..]) dnf
        groupVars = map (uniq . sort) . groupOf abs . sortOf abs
        removeTrueClauses = filter (all single)
        restoreVar = uniq . sort . map (map head)
        ref i | i > 0 =   vars !! (fromInteger (i - 1))
              | i < 0 = -(vars !! (fromInteger (-i - 1)))

-- [[A1,A2,A3..], [B1,B2,B3..],..] -> [[A1,B1,..], [A1,B2,..], .., [A2,B1,..], [A2,B2,..], ..]
multiply :: [[a]] -> [[a]]
multiply [] = []
multiply [xs] = map (\x -> [x]) xs
multiply (xs:xss) = concatMap (\x -> map (x:) $ multiply xss) xs

sandwich (a:_) []      = [a]
sandwich (a:as) (b:bs) = a:b: sandwich as bs

sortOf :: (Ord a, Ord b) => (a -> b) -> [a] -> [a]
sortOf f = map snd . sort . map (\x -> (f x, x))

groupOf f = groupBy (\a b -> f a == f b)

pack n [] = []
pack n xs = (take n xs) : pack n (drop n xs)

single [x] = True
single _   = False

-- Eliminates adjacent same elements, same as unix command.
uniq :: Eq a => [a] -> [a]
uniq = map head . group