module Main where import Control.Applicative ((<|>)) import Data.Function ((&)) import Data.Vector ((!)) 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 qualified Data.Vector import qualified Data.Vector.Unboxed import qualified Data.Vector.Unboxed.Mutable import qualified Data.STRef 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 Grid = Data.Vector.Vector Cell type CellIxs = [Int] isPossible :: Cell -> Bool isPossible (Possible _) = True isPossible _ = False readGrid :: String -> Maybe Grid readGrid s | length s == 81 = Data.Vector.fromList <$> traverse readCell 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 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 g = g Data.Vector.// [(i, c)] showGrid :: Grid -> String showGrid grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs where showCell (Fixed x) = show . Data.Bits.countTrailingZeros $ x showCell _ = "." showGridWithPossibilities :: Grid -> String showGridWithPossibilities grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs 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] ++ "]" cellIndicesList :: [Cell] -> [Data.Word.Word16] cellIndicesList cells = Data.Vector.Unboxed.toList $ Control.Monad.ST.runST $ do vec <- Data.Vector.Unboxed.Mutable.replicate 9 Data.Bits.zeroBits ref <- Data.STRef.newSTRef (1 :: Int) Control.Monad.forM_ cells $ \cell -> do i <- Data.STRef.readSTRef ref case cell of Fixed _ -> return () Possible xs -> Control.Monad.forM_ [0..8] $ \d -> Control.Monad.when (Data.Bits.testBit xs (d+1)) $ Data.Vector.Unboxed.Mutable.unsafeModify vec (`Data.Bits.setBit` i) d Data.STRef.writeSTRef ref (i+1) Data.Vector.Unboxed.unsafeFreeze vec exclusivePossibilities :: [Cell] -> [Data.Word.Word16] exclusivePossibilities row = row & cellIndicesList & zip [1..9] & filter (\(_, xs) -> let p = Data.Bits.popCount xs in p > 0 && p < 4) & Data.List.foldl' (\acc (x, is) -> Map.insertWith prepend is [x] acc) Map.empty & Map.filterWithKey (\is xs -> Data.Bits.popCount 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 pruneCells :: Grid -> CellIxs -> Maybe Grid pruneCells grid cellIxs = Control.Monad.foldM pruneCell grid cellIxs where cells = map (grid !) cellIxs exclusives = exclusivePossibilities cells allExclusives = setBits Data.Bits.zeroBits exclusives fixeds = setBits Data.Bits.zeroBits [x | Fixed x <- cells] pruneCell g i = pruneCellByFixed g (i, g ! i) >>= \g' -> pruneCellByExclusives g' (i, g' ! i) pruneCellByFixed g (_, Fixed _) = Just g pruneCellByFixed g (i, Possible xs) | xs' == xs = Just g | otherwise = flip (replaceCell i) g <$> makeCell xs' where xs' = xs Data.Bits..&. Data.Bits.complement fixeds pruneCellByExclusives g (_, Fixed _) = Just g pruneCellByExclusives g (i, Possible xs) | null exclusives = Just g | 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.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 = [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.Vector.minimumBy (compare `Data.Function.on` (possibilityCount . snd)) . Data.Vector.imapMaybe (\j cell -> if isPossible cell then Just (j, cell) else Nothing) $ grid in (replaceCell i first grid, replaceCell 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) 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'