summaryrefslogtreecommitdiff
path: root/2024/3.hs
blob: 86a5feef1c6e31e059f32291538f192a0c09495d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
{-# LANGUAGE LambdaCase #-}
import Data.List (tails, foldl')
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)

data Regex
  = L String       -- ^ literal
  | A [Regex]      -- ^ alternation (or)
  | S [Regex]      -- ^ sequence (concatenation)
  | M Regex        -- ^ many (star)
  | Set [Char]     -- ^ character class
  | Grp Int Regex  -- ^ capture group
  deriving (Show)

instance Semigroup Regex where
  a <> S bs = S (a : bs)
  a <> b = S [a, b]
instance Monoid Regex where mempty = L ""

some :: Regex -> Regex
some r = r <> M r

-- (full match, groups, rest)
match :: Regex -> String -> [(String, Map Int String, String)]
match (L r) s
  | take (length r) s == r = pure (r, mempty, drop (length r) s)
  | otherwise = []
match (A rs) s = concatMap (`match` s) rs
match (S []) s = [("", mempty, s)]
match (S (r:rs)) s = do
  (f1, m1, rest1) <- match r s
  (f2, m2, rest2) <- match (S rs) rest1
  pure (f1 ++ f2, m1 <> m2, rest2)
match (M r) s =
  case match r s of
    [] -> [("", mempty, s)]
    ms -> do (f1, m1, rest1) <- ms
             case match (M r) rest1 of
               [] -> pure (f1, m1, rest1)
               ms' -> do (f2, m2, rest2) <- ms'
                         pure (f1 ++ f2, m1 <> m2, rest2)
match (Set chs) s
  | c : s' <- s, c `elem` chs = [([c], mempty, s')]
  | otherwise = []
match (Grp i r) s = do
  (f, m, rest) <- match r s
  pure (f, Map.insert i f m, rest)

matchall :: Regex -> String -> [(String, Map Int String, String)]
matchall r s = concatMap (match r) (tails s)

main :: IO ()
main = do
  let num_reg = some (Set ['0'..'9'])
      mul_reg = L "mul(" <> Grp 1 num_reg <> L "," <> Grp 2 num_reg <> L ")"
  input <- getContents
  print $ sum [read @Int (m Map.! 1) * read (m Map.! 2) | (_, m, _) <- matchall mul_reg input]

  let instr_reg = A [mul_reg, L "do()", L "don't()"]
  let instrs = [case (,) <$> Map.lookup 1 m <*> Map.lookup 2 m of
                  Just (a, b) -> Right (read @Int a * read b)
                  Nothing -> Left (s == "do()")
               | (s, m, _) <- matchall instr_reg input]
  print $ fst $
    foldl' (\(acc, en) -> \case Right v | en -> (acc + v, en)
                                        | otherwise -> (acc, en)
                                Left en' -> (acc, en'))
           (0, True) instrs