187 lines
5.9 KiB
Haskell
187 lines
5.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.Split
|
|
import qualified Data.List
|
|
import qualified Data.Map.Strict as Map
|
|
import qualified Data.Word
|
|
import qualified Data.Bits
|
|
|
|
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..|.)
|
|
|
|
data Cell = Fixed Data.Word.Word16
|
|
| Possible Data.Word.Word16
|
|
deriving (Show, Eq)
|
|
|
|
type Row = [Cell]
|
|
type Grid = [Row]
|
|
|
|
isPossible :: Cell -> Bool
|
|
isPossible (Possible _) = True
|
|
isPossible _ = False
|
|
|
|
readGrid :: String -> Maybe Grid
|
|
readGrid s
|
|
| length s == 81 = traverse (traverse readCell) . Data.List.Split.chunksOf 9 $ s
|
|
| otherwise = Nothing
|
|
where
|
|
allBitsSet = 1022
|
|
|
|
readCell '.' = Just $ Possible allBitsSet
|
|
readCell c
|
|
| Data.Char.isDigit c && c > '0' = Just . Fixed . Data.Bits.bit . Data.Char.digitToInt $ c
|
|
| otherwise = Nothing
|
|
|
|
showGrid :: Grid -> String
|
|
showGrid = unlines . map (unwords . map showCell)
|
|
where
|
|
showCell (Fixed x) = show . Data.Bits.countTrailingZeros $ x
|
|
showCell _ = "."
|
|
|
|
showGridWithPossibilities :: Grid -> String
|
|
showGridWithPossibilities = unlines . map (unwords . map showCell)
|
|
where
|
|
showCell (Fixed x) = (show . Data.Bits.countTrailingZeros $ x) ++ " "
|
|
showCell (Possible xs) =
|
|
"[" ++ map (\i -> if Data.Bits.testBit xs i then Data.Char.intToDigit i else ' ') [1..9] ++ "]"
|
|
|
|
exclusivePossibilities :: [Cell] -> [Data.Word.Word16]
|
|
exclusivePossibilities row =
|
|
row
|
|
& zip [1..9]
|
|
& filter (isPossible . snd)
|
|
& Data.List.foldl'
|
|
(\acc ~(i, Possible xs) ->
|
|
Data.List.foldl'
|
|
(\acc' n -> if Data.Bits.testBit xs n then Map.insertWith prepend n [i] acc' else acc')
|
|
acc
|
|
[1..9])
|
|
Map.empty
|
|
& Map.filter ((< 4) . length)
|
|
& Map.foldlWithKey'(\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 ys == 1 = Just $ Fixed ys
|
|
| otherwise = Just $ Possible ys
|
|
|
|
pruneCellsByFixed :: [Cell] -> Maybe [Cell]
|
|
pruneCellsByFixed cells = traverse pruneCell cells
|
|
where
|
|
fixeds = setBits Data.Bits.zeroBits [x | Fixed x <- cells]
|
|
|
|
pruneCell (Possible xs) = makeCell (xs Data.Bits..&. Data.Bits.complement fixeds)
|
|
pruneCell x = Just x
|
|
|
|
pruneCellsByExclusives :: [Cell] -> Maybe [Cell]
|
|
pruneCellsByExclusives cells = case exclusives of
|
|
[] -> Just cells
|
|
_ -> traverse pruneCell cells
|
|
where
|
|
exclusives = exclusivePossibilities cells
|
|
allExclusives = setBits Data.Bits.zeroBits exclusives
|
|
|
|
pruneCell cell@(Fixed _) = Just cell
|
|
pruneCell cell@(Possible xs)
|
|
| intersection `elem` exclusives = makeCell intersection
|
|
| otherwise = Just cell
|
|
where
|
|
intersection = xs Data.Bits..&. allExclusives
|
|
|
|
pruneCells :: [Cell] -> Maybe [Cell]
|
|
pruneCells cells = fixM pruneCellsByFixed cells >>= fixM pruneCellsByExclusives
|
|
|
|
subGridsToRows :: Grid -> Grid
|
|
subGridsToRows =
|
|
concatMap (\rows -> let [r1, r2, r3] = map (Data.List.Split.chunksOf 3) rows
|
|
in zipWith3 (\a b c -> a ++ b ++ c) r1 r2 r3)
|
|
. Data.List.Split.chunksOf 3
|
|
|
|
pruneGrid' :: Grid -> Maybe Grid
|
|
pruneGrid' grid =
|
|
traverse pruneCells grid
|
|
>>= fmap Data.List.transpose . traverse pruneCells . Data.List.transpose
|
|
>>= fmap subGridsToRows . traverse pruneCells . subGridsToRows
|
|
|
|
pruneGrid :: Grid -> Maybe Grid
|
|
pruneGrid = fixM pruneGrid'
|
|
|
|
isGridFilled :: Grid -> Bool
|
|
isGridFilled grid = null [ () | Possible _ <- concat grid ]
|
|
|
|
isGridInvalid :: Grid -> Bool
|
|
isGridInvalid grid =
|
|
any isInvalidRow grid
|
|
|| any isInvalidRow (Data.List.transpose grid)
|
|
|| any isInvalidRow (subGridsToRows grid)
|
|
where
|
|
isInvalidRow row =
|
|
let fixeds = [x | Fixed x <- row]
|
|
emptyPossibles = [() | Possible x <- row, x == Data.Bits.zeroBits]
|
|
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@(Fixed _), rest) =
|
|
fixCell
|
|
. Data.List.minimumBy (compare `Data.Function.on` (possibilityCount . snd))
|
|
. filter (isPossible . snd)
|
|
. zip [0..]
|
|
. concat
|
|
$ grid
|
|
in (replace2D i first grid, replace2D i rest grid)
|
|
where
|
|
possibilityCount (Possible xs) = Data.Bits.popCount xs
|
|
possibilityCount (Fixed _) = 1
|
|
|
|
fixCell ~(i, Possible xs) =
|
|
let x = Data.Bits.countTrailingZeros xs
|
|
in case makeCell (Data.Bits.clearBit xs x) of
|
|
Nothing -> error "Impossible case"
|
|
Just cell -> (i, Fixed (Data.Bits.bit x), cell)
|
|
|
|
replace2D :: Int -> a -> [[a]] -> [[a]]
|
|
replace2D i v = let (x, y) = (i `quot` 9, i `mod` 9) in replace x (replace y (const v))
|
|
replace p f xs = [if i == p then f x else x | (x, i) <- zip xs [0..]]
|
|
|
|
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'
|