hasdoku/src/Sudoku.hs

223 lines
7.9 KiB
Haskell

module Main where
import Control.Applicative ((<|>))
import Data.Function ((&))
import qualified Control.Monad
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
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 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 :: [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] ]
replaceCell :: Int -> Cell -> Grid -> Grid
replaceCell i c = Data.Vector.Unboxed.modify (\v -> Data.Vector.Unboxed.Mutable.write 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 :: Grid -> CellIxs -> [Data.Word.Word16]
exclusivePossibilities grid cellIxs =
cellIxs
& map (grid !)
& 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 :: Grid -> CellIxs -> Maybe Grid
pruneCells grid cellIxs = Control.Monad.foldM pruneCell grid cellIxs
where
exclusives = exclusivePossibilities grid cellIxs
allExclusives = setBits Data.Bits.zeroBits exclusives
fixeds = setBits Data.Bits.zeroBits . filter isFixed . map (grid !) $ cellIxs
pruneCell g i =
pruneCellByFixed g (i, g ! i) >>= \g' -> pruneCellByExclusives g' (i, g' ! i)
pruneCellByFixed g (i, xs)
| isFixed xs = Just g
| xs' == xs = Just g
| otherwise = flip (replaceCell i) g <$> makeCell xs'
where
xs' = xs Data.Bits..&. Data.Bits.complement fixeds
pruneCellByExclusives g (i, xs)
| isFixed xs = Just g
| null exclusives = Just g
| makePossible intersection == xs = Just g
| intersection `elem` exclusives =
flip (replaceCell i) g <$> makeCell intersection
| otherwise = Just g
where
intersection = xs Data.Bits..&. allExclusives
pruneGrid' :: Grid -> Maybe Grid
pruneGrid' grid =
Control.Monad.foldM pruneCells grid allRowIxs
>>= flip (Control.Monad.foldM pruneCells) allColIxs
>>= flip (Control.Monad.foldM pruneCells) allSubGridIxs
pruneGrid :: Grid -> Maybe Grid
pruneGrid = fixM pruneGrid'
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
inputs <- lines <$> getContents
Control.Monad.forM_ inputs $ \input ->
case readGrid input of
Nothing -> putStrLn "Invalid input"
Just grid -> case solve grid of
Nothing -> putStrLn "No solution found"
Just grid' -> putStrLn $ showGrid grid'