diff --git a/app/Day6.hs b/app/Day6.hs index 838058b..f186eda 100644 --- a/app/Day6.hs +++ b/app/Day6.hs @@ -1,5 +1,6 @@ module Main where +import Data.Char (digitToInt) import Data.Either (lefts, rights) import Data.IntMap (IntMap) import qualified Data.IntMap as M @@ -14,16 +15,13 @@ main = do part1 :: String -> Int part1 = M.foldl' (+) 0 . parse - -data Oper = Mult | Add - deriving (Show) - --- Parse each space-delimited value as either a number or an --- operator. Then combine the operator and the operands into a --- calculation. -parse :: String -> IntMap Int -parse = M.map complete . foldl' insertPartial mempty . concat . map parseLine . lines where + -- Parse each space-delimited value as either a number or an + -- operator. Then combine the operator and the operands into a + -- calculation. + parse :: String -> IntMap Int + parse = M.map complete . foldl' insertPartial mempty . concat . map parseLine . lines + parseLine :: String -> [(Int, Either Oper Int)] parseLine = zip [0 ..] . map parseWord . words @@ -40,14 +38,95 @@ parse = M.map complete . foldl' insertPartial mempty . concat . map parseLine . (nums, Mult) -> product nums (nums, Add) -> sum nums +data Oper = Mult | Add + deriving (Show) + part2 :: String -> Int -part2 = error "Not implemented" +part2 = (\g -> foldl' (+) 0 $ map (calculateGroup g) [0 .. M.size (columns g) - 1]) . newGrid + +-- Extract the character in the grid on row, group and column. +-- +-- Return Nothing if the char is ' '. +idx :: Grid -> Int -> Int -> Int -> Maybe Char +idx (Grid cols rows cs) row col n + | col >= M.size cols = error $ "idx: requested a too large number group index (" ++ show col ++ " >= " ++ show cols ++ ")" + | row >= rows = error $ "idx: requested a too large row index (" ++ show row ++ " >= " ++ show rows ++ ")" + | n >= size (cols M.! col) = error $ "idx: requested a too large column index (" ++ show n ++ " >= " ++ show (size (cols M.! col)) ++ ")" + | otherwise = case cs !! row !! i of + ' ' -> Nothing + c -> Just c + where + i = n + (sum . map ((+ 1) . (size . (cols M.!))) $ [0 .. col - 1]) + +digit :: Grid -> Int -> Int -> Int -> Maybe Int +digit g r c = fmap digitToInt . idx g r c + +operator :: Grid -> Int -> Oper +operator g col = case idx g (rows g - 1) col 0 of + Just '+' -> Add + Just '*' -> Mult + Just ch -> error $ "operator: Expected '+' or '*', got " <> show ch + Nothing -> error $ "operator: Expected '+' or '*', got nothing" + +-- Calculate a column by traversing all rows at that column inside that group. +-- +-- Calculation is performed by first multiplying the accumulated value +-- by 10 and then adding the new digit. If no value exists in that +-- slot we don't modify the accumulated value. +calculateColumn :: Grid -> Int -> Int -> Int +calculateColumn g c n = foldl' (\acc r -> maybe acc ((+) (acc * 10)) $ digit g r c n) 0 [0 .. rows g - 2] + +-- Calculates each column and folds them using the operator at the last line. +calculateGroup :: Grid -> Int -> Int +calculateGroup g c = foldl' (op $ oper column) (identity column) $ map (calculateColumn g c) [0 .. (size column) - 1] + where + column = columns g M.! c + +op :: Oper -> (Int -> Int -> Int) +op Mult = (*) +op Add = (+) + +data Grid = Grid + { columns :: IntMap Column, + rows :: Int, + cells :: [[Char]] + } + deriving (Show) + +data Column = Column + { size :: Int, + oper :: Oper, + identity :: Int + } + deriving (Show) + +-- | Construct a grid based on the last line. +-- +-- Using the last line we can get the number of digits in each column. +newGrid :: String -> Grid +newGrid = go . lines + where + go :: [String] -> Grid + go ls = + let columns = M.fromList $ foldl' go' [] (last ls) + in Grid columns (length ls) ls + + go' :: [(Int, Column)] -> Char -> [(Int, Column)] + go' [] ch = pure . (0,) . (\(oper, idty) -> Column 1 oper idty) $ parseOper ch + go' ((n, c) : lengths) ch + | ch == ' ' = (n, c {size = size c + 1}) : lengths + | otherwise = (n + 1, (\(oper, idty) -> Column 1 oper idty) (parseOper ch)) : (n, c {size = size c - 1}) : lengths + + parseOper :: Char -> (Oper, Int) + parseOper '*' = (Mult, 1) + parseOper '+' = (Add, 0) + parseOper c = error $ "parseOper: expects a '+' or '*', got " <> show c testInput :: String testInput = unlines - [ "123 328 51 64", - " 45 64 387 23", - " 6 98 215 314", - "* + * + " + [ "123 328 51 64 1", + " 45 64 387 23 2", + " 6 98 215 314 3", + "* + * + +" ]