{-# Language ImportQualifiedPost, QuasiQuotes, BlockArguments #-}
{-|
Module      : Main
Description : Day 12 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/12>

Search around a cave visiting some caves more than others.

This solution makes the observation that we can optimize
away all the big caves. Big caves can never be connected
to other big caves or we'd have infinite cycles, and we
don't need to track anything about visiting a big cave.

-}
module Main (main) where

import Advent.Format (format)
import Advent.Memo (memo3)
import Advent.SmallSet qualified as SmallSet
import Control.Monad.Trans.State.Strict
import Data.Char (isUpper)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.List (mapAccumL)
import Data.Map (Map)
import Data.Map qualified as Map

-- | >>> :main
-- 3761
-- 99138
main :: IO ()
IO ()
main =
 do IntMap (IntMap Int)
inp <- IntMap [Int] -> IntMap (IntMap Int)
compress (IntMap [Int] -> IntMap (IntMap Int))
-> ([(String, String)] -> IntMap [Int])
-> [(String, String)]
-> IntMap (IntMap Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Int)] -> IntMap [Int]
toAdj ([(Int, Int)] -> IntMap [Int])
-> ([(String, String)] -> [(Int, Int)])
-> [(String, String)]
-> IntMap [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(String, String)] -> [(Int, Int)]
label ([(String, String)] -> IntMap (IntMap Int))
-> IO [(String, String)] -> IO (IntMap (IntMap Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2021 12 (%s-%s%n)*|]
    Int -> IO ()
forall a. Show a => a -> IO ()
print (IntMap (IntMap Int) -> Bool -> Int
start IntMap (IntMap Int)
inp Bool
False)
    Int -> IO ()
forall a. Show a => a -> IO ()
print (IntMap (IntMap Int) -> Bool -> Int
start IntMap (IntMap Int)
inp Bool
True)

-- | Compute directed edge map from a list of undirected edges.
toAdj :: [(Int, Int)] -> IntMap [Int]
toAdj :: [(Int, Int)] -> IntMap [Int]
toAdj [(Int, Int)]
inp = ([Int] -> [Int] -> [Int]) -> [(Int, [Int])] -> IntMap [Int]
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
(++)
  [(Int
x,[Int
y]) | (Int
a,Int
b) <- [(Int, Int)]
inp, (Int
x,Int
y) <- [(Int
a,Int
b),(Int
b,Int
a)], Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0, Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1]

-- | Compute direct paths through a big cave to the next small cave.
compress :: IntMap [Int] -> IntMap (IntMap Int)
compress :: IntMap [Int] -> IntMap (IntMap Int)
compress IntMap [Int]
long = (Int -> IntMap Int -> Bool)
-> IntMap (IntMap Int) -> IntMap (IntMap Int)
forall a. (Int -> a -> Bool) -> IntMap a -> IntMap a
IntMap.filterWithKey (\Int
k IntMap Int
_ -> Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) IntMap (IntMap Int)
short
  where
    short :: IntMap (IntMap Int)
short = [Int] -> IntMap Int
shorten ([Int] -> IntMap Int) -> IntMap [Int] -> IntMap (IntMap Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap [Int]
long
    shorten :: [Int] -> IntMap Int
shorten [Int]
xs = (Int -> Int -> Int) -> [IntMap Int] -> IntMap Int
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IntMap.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)
      [if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IntMap.singleton Int
x Int
1 else IntMap (IntMap Int)
short IntMap (IntMap Int) -> Int -> IntMap Int
forall a. IntMap a -> Int -> a
IntMap.! Int
x | Int
x <- [Int]
xs]

-- | Search the cave exploration given the directed edges and a
-- flag if we're allowed to visit a small cave an extra time.
start :: IntMap (IntMap Int) -> Bool -> Int
start :: IntMap (IntMap Int) -> Bool -> Int
start IntMap (IntMap Int)
paths = Int -> SmallSet -> Bool -> Int
go Int
0 SmallSet
SmallSet.empty
  where
    go :: Int -> SmallSet -> Bool -> Int
go = (Int -> SmallSet -> Bool -> Int) -> Int -> SmallSet -> Bool -> Int
forall r s t a.
(HasTrie r, HasTrie s, HasTrie t) =>
(r -> s -> t -> a) -> r -> s -> t -> a
memo3 \Int
here SmallSet
seen Bool
extra ->
      let
        f :: Int -> Int
f Int
next
          | Int
next Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Int
1
          | Bool -> Bool
not (Int -> SmallSet -> Bool
SmallSet.member Int
next SmallSet
seen) = Int -> SmallSet -> Bool -> Int
go Int
next (Int -> SmallSet -> SmallSet
SmallSet.insert Int
next SmallSet
seen) Bool
extra
          | Bool
extra     = Int -> SmallSet -> Bool -> Int
go Int
next SmallSet
seen Bool
False
          | Bool
otherwise = Int
0
      in [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
f Int
k | (Int
k,Int
v) <- IntMap Int -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
IntMap.toList (IntMap (IntMap Int)
paths IntMap (IntMap Int) -> Int -> IntMap Int
forall a. IntMap a -> Int -> a
IntMap.! Int
here)]

-- | Map all the cave names to integers. Use negative integers for big caves.
-- @start@ always gets assigned @0@ and @end@ gets @1@
label :: [(String, String)] -> [(Int,Int)]
label :: [(String, String)] -> [(Int, Int)]
label [(String, String)]
xs =
  State (Map String Int) [(Int, Int)]
-> Map String Int -> [(Int, Int)]
forall s a. State s a -> s -> a
evalState
    (((String, String) -> StateT (Map String Int) Identity (Int, Int))
-> [(String, String)] -> State (Map String Int) [(Int, Int)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((String -> StateT (Map String Int) Identity Int)
-> (String, String) -> StateT (Map String Int) Identity (Int, Int)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (a, a) -> f (b, b)
both ((Map String Int -> (Int, Map String Int))
-> StateT (Map String Int) Identity Int
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((Map String Int -> (Int, Map String Int))
 -> StateT (Map String Int) Identity Int)
-> (String -> Map String Int -> (Int, Map String Int))
-> String
-> StateT (Map String Int) Identity Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Map String Int -> (Int, Map String Int)
label1)) [(String, String)]
xs)
    ([(String, Int)] -> Map String Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(String
"start",Int
0), (String
"end",Int
1)])

-- | Assigns a unique integer label for each cave name. Big cave names
-- are assigned negative integers.
label1 :: String -> Map String Int -> (Int, Map String Int)
label1 :: String -> Map String Int -> (Int, Map String Int)
label1 String
x Map String Int
m =
  case String -> Map String Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
x Map String Int
m of
    Just Int
i  -> (Int
i, Map String Int
m)
    Maybe Int
Nothing -> (Int
i, String -> Int -> Map String Int -> Map String Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
x Int
i Map String Int
m)
      where i :: Int
i = if Char -> Bool
isUpper (String -> Char
forall a. HasCallStack => [a] -> a
head String
x) then -Map String Int -> Int
forall k a. Map k a -> Int
Map.size Map String Int
m else Map String Int -> Int
forall k a. Map k a -> Int
Map.size Map String Int
m

-- | Traverse over the first and second components of a pair.
both :: Applicative f => (a -> f b) -> (a,a) -> f (b,b)
both :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (a, a) -> f (b, b)
both a -> f b
f (a
x,a
y) = (,) (b -> b -> (b, b)) -> f b -> f (b -> (b, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
x f (b -> (b, b)) -> f b -> f (b, b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> f b
f a
y