135 lines
4.4 KiB
Haskell
135 lines
4.4 KiB
Haskell
module Main where
|
|
|
|
import Data.Char (digitToInt)
|
|
import Data.Either (lefts, rights)
|
|
import Data.IntMap (IntMap)
|
|
import qualified Data.IntMap as M
|
|
import Data.List (foldl')
|
|
import System.Environment (getArgs)
|
|
|
|
-- TODO: solve part 2 by transposing the matrix before parsing the
|
|
-- numbers.
|
|
|
|
main :: IO ()
|
|
main = do
|
|
input <- readFile . head =<< getArgs
|
|
putStrLn $ unwords ["Part 1:", show $ part1 input]
|
|
putStrLn $ unwords ["Part 2:", show $ part2 input]
|
|
|
|
part1 :: String -> Int
|
|
part1 = M.foldl' (+) 0 . parse
|
|
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
|
|
|
|
parseWord :: String -> Either Oper Int
|
|
parseWord "*" = Left Mult
|
|
parseWord "+" = Left Add
|
|
parseWord w = pure $ read w
|
|
|
|
insertPartial :: IntMap [Either Oper Int] -> (Int, Either Oper Int) -> IntMap [Either Oper Int]
|
|
insertPartial m (column, oper) = M.insertWith ((<>)) column [oper] m
|
|
|
|
complete :: [Either Oper Int] -> Int
|
|
complete m = case (rights m, head (lefts m)) of
|
|
(nums, Mult) -> product nums
|
|
(nums, Add) -> sum nums
|
|
|
|
data Oper = Mult | Add
|
|
deriving (Show)
|
|
|
|
part2 :: String -> Int
|
|
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 1",
|
|
" 45 64 387 23 2",
|
|
" 6 98 215 314 3",
|
|
"* + * + +"
|
|
]
|