143 lines
5.8 KiB
Haskell
143 lines
5.8 KiB
Haskell
import qualified Data.Map as M
|
|
import qualified Data.Set as S
|
|
import Data.List.Split (splitOn)
|
|
import Debug.Trace (traceShowId, traceShow, trace)
|
|
import Data.Maybe (fromMaybe)
|
|
|
|
import Bfs
|
|
import Data.List (maximumBy, permutations, delete, singleton)
|
|
import Data.Function (on)
|
|
import Control.Monad.Trans.State (State, get, put, evalState)
|
|
|
|
data Valve = V { flow' :: Int, neighs' :: [String] } deriving Show
|
|
type ZMap = M.Map String Valve
|
|
|
|
data NzValve = NV { flow :: Int, dists :: M.Map String Int} deriving Show
|
|
type NzMap = M.Map String NzValve
|
|
|
|
parseValve :: String -> (String, Valve)
|
|
parseValve xs = (name, V flow tunnels) where
|
|
name = words xs !! 1
|
|
flow = read $ last $ splitOn "=" (head $ splitOn ";" xs)
|
|
valves = splitOn ", " $ last (splitOn "tunnels lead to valves " xs)
|
|
singleValve = splitOn "tunnel leads to valve " xs
|
|
tunnels = if length singleValve == 1 then valves else [last singleValve]
|
|
|
|
search' :: ZMap -> String -> String -> Int
|
|
search' scan from to = dists M.! to where
|
|
dists = bfs f [(from,0)] M.empty
|
|
f x = neighs' (scan M.! x)
|
|
|
|
|
|
prepass :: ZMap -> NzMap
|
|
prepass xs = M.fromList $ concatMap f $ M.toList xs where
|
|
f (name, v@(V flow _)) = [(name, NV flow (M.fromList $ dists name)) | name == "AA" || flow > 0]
|
|
dists name = map (\target -> (target, search' xs name target)) (filter (/= name) allValves)
|
|
allValves = map fst . filter (\(name, V flow _) -> name == "AA" || flow > 0) $ M.toList xs
|
|
|
|
|
|
|
|
|
|
type Memo = M.Map (S.Set String, String, Int) Int
|
|
-- search :: Map -> S.Set String -> String -> Int -> Memo -> (Int, Memo)
|
|
-- search scan active location n memo
|
|
-- | n <= 0 = (0, M.empty)
|
|
-- | otherwise = (releasedPressure + fst best, snd best) where
|
|
-- valve = scan M.! location
|
|
-- openValve = [(S.insert location active, location, n-1) | S.notMember location active && flow valve > 0]
|
|
-- releasedPressure = sum $ map (flow . (scan M.!)) (S.toList active)
|
|
-- nop = [(active, location, n-1) | S.size active == M.size scan]
|
|
-- move = map (\(name, dist) -> (active, name, n-dist)) (dists $ scan M.! location)
|
|
-- actions = openValve ++ move ++ nop
|
|
-- trace x = if n>=15 then traceShow (fst x) x else x
|
|
-- best = trace $ maximumBy (compare `on` fst) $ map (\(a,l,n) -> search scan a l n memo) actions
|
|
-- -- best = maybe best' (\b -> (b,memo)) (M.lookup (active, location, n) memo)
|
|
-- -- best' = foldr f (0, memo) actions where
|
|
-- -- f (a,l, newN) (i,m) = let (result, memo') = search scan a l newN m; b = max i result in
|
|
-- -- trace (b, M.insertWith max (a,l, newN) b $ M.unionWith max m memo')
|
|
|
|
perms :: Int -> NzMap -> [[String]]
|
|
perms limit m = perms' "AA" limit elems where
|
|
elems = filter (/= "AA") $ M.keys m
|
|
perms' :: String -> Int -> [String] -> [[String]]
|
|
perms' current remaining [] = [[]]
|
|
perms' current remaining as = do
|
|
a <- as
|
|
let distance = dists (m M.! current) M.! a
|
|
[a] : if distance < remaining then do
|
|
let l = delete a as
|
|
ls <- perms' a (remaining-distance) l
|
|
pure $ a:ls
|
|
else pure [a]
|
|
|
|
execPerm :: NzMap -> [String] -> Int -> String -> Int -> Int
|
|
execPerm nzmap [] flowrate location remaining = remaining * flowrate
|
|
execPerm nzmap (target:qs) flowrate location remaining
|
|
| remaining < 0 = 0
|
|
| distance >= remaining = remaining * flowrate
|
|
| otherwise = (distance * flowrate) + execPerm nzmap qs flowrate' target (remaining - distance)
|
|
where
|
|
valve = nzmap M.! location
|
|
distance = 1 + dists valve M.! target
|
|
flowrate' = flowrate + flow (nzmap M.! target)
|
|
|
|
|
|
traceMap :: (Show a, Show b) => (a -> a -> Bool) -> (a -> b) -> [a] -> [b]
|
|
traceMap cond f [] = []
|
|
traceMap cond f [x] = [f x]
|
|
traceMap cond f (x:y:xs)
|
|
| cond x y = let result = f x in traceShow (result, x) (result: traceMap cond f (y:xs))
|
|
| otherwise = f x : traceMap cond f (y:xs)
|
|
|
|
type EleMemo = M.Map (S.Set String) Int
|
|
eleLookup :: NzMap -> S.Set String -> State EleMemo Int
|
|
eleLookup nzmap permutation = do
|
|
mmap <- get
|
|
case M.lookup permutation mmap of
|
|
Just x -> pure x
|
|
Nothing -> do
|
|
let valves = M.fromList $ filter (\(name, valve) -> S.notMember name permutation) (M.toList nzmap)
|
|
let elePerms = traceShow ("elephant", permutation) $ perms 26 valves
|
|
let exec p = execPerm nzmap p 0 "AA" 26
|
|
let flow = maximum $ map exec elePerms
|
|
put $ M.insert permutation flow mmap
|
|
pure flow
|
|
|
|
part2 :: NzMap -> [String] -> State EleMemo Int
|
|
part2 nzmap permutation = do
|
|
let me = execPerm nzmap permutation 0 "AA" 26
|
|
elephant <- eleLookup nzmap (S.fromList permutation)
|
|
pure $ me + elephant
|
|
-- case elephant of
|
|
-- Just ele -> pure $ me + ele
|
|
-- Nothing -> do
|
|
-- let valves = M.fromList $ filter (\(name, valve) -> S.notMember name permset) (M.toList nzmap)
|
|
-- let elePerms = traceShow ("elephant", permutation) $ perms 26 valves
|
|
-- let exec p = execPerm nzmap p 0 "AA" 26
|
|
-- let ele = maximum $ map exec elePerms
|
|
-- put $ M.insert permset ele state
|
|
-- pure $ me + ele
|
|
|
|
main :: IO ()
|
|
main = do
|
|
x <- prepass . M.fromList . map parseValve . lines <$> getContents
|
|
mapM_ print x
|
|
|
|
-- let all = perms 30 x
|
|
-- let exec p = execPerm x p 0 "AA" 30
|
|
-- let result = traceMap (\a b -> take 2 a /= take 2 b) exec all
|
|
-- putStr "part 1: "
|
|
-- print $ maximum result
|
|
-- print $ fst $ search x S.empty "AA" 30 M.empty
|
|
|
|
putStr "part 2: "
|
|
-- let initialEle = maximumBy (compare `on` fst) $ map (\p -> (execPerm x p 0 "AA" 26, [S.fromList p])) (perms 26 x)
|
|
-- print initialEle
|
|
|
|
|
|
let all = perms 26 x
|
|
let action = mapM (part2 x) all
|
|
let result = evalState action M.empty
|
|
print $ maximum result
|
|
pure ()
|