hasdoku/src/Sudoku.hs

266 lines
9.7 KiB
Haskell

{-# LANGUAGE LambdaCase #-}
module Main where
import Control.Applicative ((<|>))
import Data.Function ((&))
import qualified Control.Monad
import qualified Control.Monad.ST
import qualified Data.Char
import qualified Data.Function
import qualified Data.List
import qualified Data.Map.Strict as Map
import qualified Data.Word
import qualified Data.Bits
import Data.Vector.Unboxed ((!))
import qualified Data.Vector.Unboxed.Mutable
import qualified Data.Vector.Unboxed
import qualified Data.Maybe
import Control.Parallel.Strategies (withStrategy, rdeepseq, parBuffer)
fixM :: (Eq t, Monad m) => (t -> m t) -> t -> m t
fixM f x = f x >>= \x' -> if x' == x then return x else fixM f x'
setBits :: Data.Word.Word16 -> [Data.Word.Word16] -> Data.Word.Word16
setBits = Data.List.foldl' (Data.Bits..|.)
type Cell = Data.Word.Word16
type Grid = Data.Vector.Unboxed.Vector Cell
type MGrid s = Data.Vector.Unboxed.Mutable.STVector s Cell
type CellIxs = [Int]
isPossible, isFixed :: Cell -> Bool
isPossible = flip Data.Bits.testBit 15
isFixed = not . isPossible
makePossible, makeFixed :: Data.Word.Word16 -> Cell
makePossible = flip Data.Bits.setBit 15
makeFixed = flip Data.Bits.clearBit 15
readGrid :: String -> Maybe Grid
readGrid s
| length s == 81 = Data.Vector.Unboxed.fromList <$> traverse readCell s
| otherwise = Nothing
where
allBitsSet = 1022
readCell '.' = Just $ makePossible allBitsSet
readCell c
| Data.Char.isDigit c && c > '0' = Just . Data.Bits.bit . Data.Char.digitToInt $ c
| otherwise = Nothing
fromXY :: (Int, Int) -> Int
fromXY (x, y) = x * 9 + y
allRowIxs, allColIxs, allSubGridIxs, allIxs :: [CellIxs]
allRowIxs = [getRow i | i <- [0..8]]
where getRow n = [ fromXY (n, i) | i <- [0..8] ]
allColIxs = [getCol i | i <- [0..8]]
where getCol n = [ fromXY (i, n) | i <- [0..8] ]
allSubGridIxs = [getSubGrid i | i <- [0..8]]
where getSubGrid n = let (r, c) = (n `quot` 3, n `mod` 3)
in [ fromXY (3 * r + i, 3 * c + j) | i <- [0..2], j <- [0..2] ]
allIxs = concat [allRowIxs, allColIxs, allSubGridIxs]
replaceCell :: Int -> Cell -> Grid -> Grid
replaceCell i c = Data.Vector.Unboxed.modify (\v -> Data.Vector.Unboxed.Mutable.unsafeWrite v i c)
showGrid :: Grid -> String
showGrid grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs
where
showCell xs
| isPossible xs = "."
| otherwise = show . Data.Bits.countTrailingZeros $ xs
showGridWithPossibilities :: Grid -> String
showGridWithPossibilities grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs
where
showCell xs
| isPossible xs =
"[" ++ map (\i -> if Data.Bits.testBit xs i then Data.Char.intToDigit i else ' ') [1..9] ++ "]"
| otherwise = (show . Data.Bits.countTrailingZeros $ xs) ++ " "
-- Exclusive Possibilities Accumulator
data ExPosAcc = ExPosAcc ![Int] ![Int] ![Int] ![Int] ![Int] ![Int] ![Int] ![Int] ![Int]
exPosAccEmpty :: ExPosAcc
exPosAccEmpty = ExPosAcc [] [] [] [] [] [] [] [] []
exPosAccInsert :: Int -> Int -> ExPosAcc -> ExPosAcc
exPosAccInsert 1 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc (i:v1) v2 v3 v4 v5 v6 v7 v8 v9
exPosAccInsert 2 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 (i:v2) v3 v4 v5 v6 v7 v8 v9
exPosAccInsert 3 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 (i:v3) v4 v5 v6 v7 v8 v9
exPosAccInsert 4 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 v3 (i:v4) v5 v6 v7 v8 v9
exPosAccInsert 5 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 v3 v4 (i:v5) v6 v7 v8 v9
exPosAccInsert 6 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 v3 v4 v5 (i:v6) v7 v8 v9
exPosAccInsert 7 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 v3 v4 v5 v6 (i:v7) v8 v9
exPosAccInsert 8 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 v3 v4 v5 v6 v7 (i:v8) v9
exPosAccInsert 9 i (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) = ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 (i:v9)
exPosAccInsert _ _ _ = error "Impossible"
exPosAccToList :: ExPosAcc -> [(Int, [Int])]
exPosAccToList (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) =
[(1, v1), (2, v2), (3, v3), (4, v4), (5, v5), (6, v6), (7, v7), (8, v8), (9, v9)]
exclusivePossibilities :: [Cell] -> [Data.Word.Word16]
exclusivePossibilities cells =
cells
& zip ([1..9] :: [Int])
& filter (isPossible . snd)
& Data.List.foldl'
(\acc (i, xs) ->
Data.List.foldl'
(\acc' n -> if Data.Bits.testBit xs n then exPosAccInsert n i acc' else acc')
acc
[1..9])
exPosAccEmpty
& exPosAccToList
& filter ((< 4) . length . snd)
& Data.List.foldl' (\acc (x, is) -> Map.insertWith prepend is [x] acc) Map.empty
& Map.filterWithKey (\is xs -> length is == length xs)
& Map.elems
& map (Data.List.foldl' Data.Bits.setBit Data.Bits.zeroBits)
where
prepend ~[y] ys = y:ys
makeCell :: Data.Word.Word16 -> Maybe Cell
makeCell ys
| ys == Data.Bits.zeroBits = Nothing
| Data.Bits.popCount (makeFixed ys) == 1 = Just $ makeFixed ys
| otherwise = Just $ makePossible ys
pruneCells :: MGrid s -> CellIxs -> Control.Monad.ST.ST s (Maybe Bool)
pruneCells grid cellIxs = do
cells <- traverse (Data.Vector.Unboxed.Mutable.unsafeRead grid) cellIxs
pruneCells' grid cells cellIxs
pruneCells' :: MGrid s -> [Cell] -> CellIxs -> Control.Monad.ST.ST s (Maybe Bool)
pruneCells' grid cells = Control.Monad.foldM pruneCell' (Just False)
where
pruneCell' Nothing _ = return Nothing
pruneCell' (Just changed) i = pruneCell i >>= \case
Nothing -> return Nothing
Just changed' -> return $ Just (changed || changed')
exclusives = exclusivePossibilities cells
allExclusives = setBits Data.Bits.zeroBits exclusives
fixeds = setBits Data.Bits.zeroBits . filter isFixed $ cells
pruneCell i = do
cell <- Data.Vector.Unboxed.Mutable.unsafeRead grid i
pruneCellByFixed (i, cell) >>= \case
Nothing -> return Nothing
Just changed -> do
cell' <- Data.Vector.Unboxed.Mutable.unsafeRead grid i
pruneCellByExclusives (i, cell') >>= \case
Nothing -> return Nothing
Just changed' -> return $ Just (changed || changed')
pruneCellByFixed (i, xs)
| isFixed xs = return $ Just False
| xs' == xs = return $ Just False
| otherwise = case makeCell xs' of
Nothing -> return Nothing
Just cell' -> do
Data.Vector.Unboxed.Mutable.unsafeWrite grid i cell'
return $ Just True
where
xs' = xs Data.Bits..&. Data.Bits.complement fixeds
pruneCellByExclusives (i, xs)
| isFixed xs = return $ Just False
| null exclusives = return $ Just False
| makePossible intersection == xs = return $ Just False
| intersection `elem` exclusives = case makeCell intersection of
Nothing -> return Nothing
Just cell' -> do
Data.Vector.Unboxed.Mutable.unsafeWrite grid i cell'
return $ Just True
| otherwise = return $ Just False
where
intersection = xs Data.Bits..&. allExclusives
pruneGrid' :: MGrid s -> Control.Monad.ST.ST s (Maybe Bool)
pruneGrid' grid = Control.Monad.foldM pruneCells' (Just False) allIxs
where
pruneCells' Nothing _ = return Nothing
pruneCells' (Just changed) ixs = pruneCells grid ixs >>= \case
Nothing -> return Nothing
Just changed' -> return $ Just (changed || changed')
pruneGrid :: Grid -> Maybe Grid
pruneGrid grid = Control.Monad.ST.runST $ do
mGrid <- Data.Vector.Unboxed.unsafeThaw grid
fix mGrid
where
fix mg = pruneGrid' mg >>= \case
Nothing -> return Nothing
Just True -> fix mg
Just False -> Just <$> Data.Vector.Unboxed.unsafeFreeze mg
isGridFilled :: Grid -> Bool
isGridFilled = not . Data.Vector.Unboxed.any isPossible
isGridInvalid :: Grid -> Bool
isGridInvalid grid =
any isInvalidRow (map (map (grid !)) allRowIxs)
|| any isInvalidRow (map (map (grid !)) allColIxs)
|| any isInvalidRow (map (map (grid !)) allSubGridIxs)
where
isInvalidRow row =
let fixeds = filter isFixed row
emptyPossibles = filter (== 32768) . filter isPossible $ row
in hasDups fixeds || not (null emptyPossibles)
hasDups l = hasDups' l []
hasDups' [] _ = False
hasDups' (y:ys) xs
| y `elem` xs = True
| otherwise = hasDups' ys (y:xs)
nextGrids :: Grid -> (Grid, Grid)
nextGrids grid =
let (i, first, rest) =
fixCell
. Data.Vector.Unboxed.minimumBy (compare `Data.Function.on` (possibilityCount . snd))
. Data.Vector.Unboxed.imapMaybe (\j cell -> if isPossible cell then Just (j, cell) else Nothing)
$ grid
in (replaceCell i first grid, replaceCell i rest grid)
where
possibilityCount xs
| isPossible xs = Data.Bits.popCount xs - 1
| otherwise = 1
fixCell (i, xs) =
let x = Data.Bits.countTrailingZeros xs
in case makeCell (Data.Bits.clearBit xs x) of
Nothing -> error "Impossible case"
Just cell -> (i, Data.Bits.bit x, cell)
solve :: Grid -> Maybe Grid
solve grid = pruneGrid grid >>= solve'
where
solve' g
| isGridInvalid g = Nothing
| isGridFilled g = Just g
| otherwise =
let (grid1, grid2) = nextGrids g
in solve grid1 <|> solve grid2
main :: IO ()
main = do
grids <- lines <$> getContents
let solutions = parMap readAndSolve grids
putStrLn $
show (length $ filter Data.Maybe.isJust solutions) ++ "/" ++ show (length grids) ++ " solved"
where
readAndSolve grid = case readGrid grid of
Nothing -> Nothing
Just b -> solve b
chunkSize = 1000
parMap f = withStrategy (parBuffer chunkSize rdeepseq) . map f