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 qualified Data.Vector import Data.Vector ((!)) 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] ++ "]" 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 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'