126 lines
4.0 KiB
Haskell
126 lines
4.0 KiB
Haskell
module Day3 where
|
||
|
||
import Data.Char (isDigit)
|
||
import Data.Matrix (Matrix(..))
|
||
import Data.Maybe
|
||
import qualified Data.Char as C
|
||
import qualified Data.Matrix as M
|
||
import qualified Data.Vector as V
|
||
import qualified Data.Bifunctor as BF
|
||
|
||
-- import Debug.Trace
|
||
|
||
main :: IO ()
|
||
main = do
|
||
m <- M.fromLists . lines <$> readFile "inputs/day3.input"
|
||
putStrLn "Part 1 result:"
|
||
print $ solvePart1 m (symbolZone m)
|
||
putStrLn "Part 2 result:"
|
||
print $ solvePart2 m
|
||
|
||
newtype Symbol = Symbol Bool
|
||
deriving Eq
|
||
|
||
instance Show Symbol where
|
||
show (Symbol True) = "X"
|
||
show (Symbol False) = " "
|
||
|
||
solvePart2 :: Matrix Char -> Int
|
||
solvePart2 = sum . cogs
|
||
|
||
countParts :: (Int, Int) -> Matrix Char -> Int
|
||
countParts (x,y) = sum
|
||
. map tallyRow
|
||
. fmap (fmap justDigits)
|
||
. flip neighborList (x,y)
|
||
|
||
justDigits :: Maybe Char -> Maybe Char
|
||
justDigits (Just d) | isDigit d = Just d
|
||
justDigits _ = Nothing
|
||
|
||
tallyRow :: [Maybe Char] -> Int
|
||
tallyRow [ Just _ , Just _ , Just _ ] = 1
|
||
tallyRow [ Nothing , Just _ , Just _ ] = 1
|
||
tallyRow [ Just _ , Nothing , Just _ ] = 2
|
||
tallyRow [ Nothing , Nothing , Just _ ] = 1
|
||
tallyRow [ Just _ , Just _ , Nothing ] = 1
|
||
tallyRow [ Nothing , Just _ , Nothing ] = 1
|
||
tallyRow [ Just _ , Nothing , Nothing ] = 1
|
||
tallyRow [ Nothing , Nothing , Nothing ] = 0
|
||
tallyRow _ = error "This should never happen"
|
||
|
||
neighborList :: Matrix a -> (Int, Int) -> [[Maybe a]]
|
||
neighborList m (x,y) =
|
||
[ [ M.safeGet (x-1) (y-1) m , M.safeGet (x-1) y m , M.safeGet (x-1) (y+1) m ]
|
||
, [ M.safeGet x (y-1) m , M.safeGet x y m , M.safeGet x (y+1) m ]
|
||
, [ M.safeGet (x+1) (y-1) m , M.safeGet (x+1) y m , M.safeGet (x+1) (y+1) m ]
|
||
]
|
||
|
||
cogs :: Matrix Char -> Matrix Int
|
||
cogs m = M.mapPos go m
|
||
where
|
||
go p e | e == '*' =
|
||
if (== 2) $ countParts p m
|
||
then product $ numbersAround p m
|
||
else 0
|
||
go _ _ = 0
|
||
|
||
numbersAround :: Read b => (Int, Int) -> Matrix Char -> [b]
|
||
numbersAround (x,y) m =
|
||
concatMap (map (read . V.toList . snd) . filter (\ (p,_) -> V.any (\ y' -> y' >= y-1 && y' <= y+1) p)
|
||
. filter (V.all isDigit . snd)
|
||
. collectDigitAndCoord) rows
|
||
where
|
||
rows = [M.getRow (x-1) m,M.getRow x m,M.getRow (x+1) m]
|
||
collectDigitAndCoord v =
|
||
map (BF.first (V.map (+ 1)) . V.unzip) $ V.groupBy (\ (_,v1) (_,v2) -> isDigit v1 && isDigit v2) $ V.indexed v
|
||
|
||
solvePart1 :: Matrix Char -> Matrix Symbol -> Int
|
||
solvePart1 m = sum . map read . words
|
||
. M.toList
|
||
. mask m
|
||
. go' . go'
|
||
where
|
||
mask = M.elementwise (\ em es -> if symbolAndDigit (em,es) then em else ' ')
|
||
go' s = M.mapPos (go s) s
|
||
go _ _ (Symbol True) = Symbol True
|
||
go s' (x,y) _ = Symbol $
|
||
any symbolAndDigit (zip (neighborList' (x,y) m) (neighborList' (x,y) s'))
|
||
neighborList' (x,y) m' =
|
||
catMaybes [ M.safeGet x (y-1) m'
|
||
, M.safeGet x (y+1) m'
|
||
]
|
||
symbolAndDigit (v, Symbol True) | C.isDigit v = True
|
||
symbolAndDigit _ = False
|
||
_noSymbolAndDigit (v, Symbol False) | C.isDigit v = True
|
||
_noSymbolAndDigit _ = False
|
||
|
||
symbolZone :: Matrix Char -> Matrix Symbol
|
||
symbolZone m = M.mapPos go m
|
||
where
|
||
go (x,y) _ = Symbol $ any isSymbol (neighborList' (x,y))
|
||
neighborList' :: (Int,Int) -> [Char]
|
||
neighborList' (x,y) = catMaybes [ M.safeGet (x-1) y m
|
||
, M.safeGet x y m
|
||
, M.safeGet (x+1) y m
|
||
, M.safeGet (x-1) (y+1) m
|
||
, M.safeGet x (y+1) m
|
||
, M.safeGet (x+1) (y+1) m
|
||
, M.safeGet (x-1) (y-1) m
|
||
, M.safeGet x (y-1) m
|
||
, M.safeGet (x+1) (y-1) m
|
||
]
|
||
|
||
isSymbol :: Char -> Bool
|
||
isSymbol '#' = True
|
||
isSymbol '$' = True
|
||
isSymbol '%' = True
|
||
isSymbol '&' = True
|
||
isSymbol '*' = True
|
||
isSymbol '+' = True
|
||
isSymbol '-' = True
|
||
isSymbol '/' = True
|
||
isSymbol '=' = True
|
||
isSymbol '@' = True
|
||
isSymbol _ = False
|