72 lines
2.5 KiB
Haskell
72 lines
2.5 KiB
Haskell
{-# LANGUAGE TupleSections #-}
|
|
|
|
import Common
|
|
import AStar
|
|
import Data.List.Split (splitOn)
|
|
import Data.Maybe (catMaybes)
|
|
import Data.List (find, maximumBy)
|
|
import GHC.Data.Maybe (fromJust)
|
|
import GHC.Utils.Misc (sndOf3)
|
|
|
|
data Valve = Valve {
|
|
valveName :: String,
|
|
valveFlowRate :: Int,
|
|
valveLeadsTo :: [String],
|
|
valveOpen :: Bool
|
|
}
|
|
deriving (Show, Eq)
|
|
|
|
-- disgusting, but works
|
|
parse :: String -> Valve
|
|
parse s = Valve name flowRate leadsTo False
|
|
where
|
|
leadsTo = map trim $ splitOn "," $ unwords $ drop 9 w
|
|
name = w !! 1
|
|
flowRate = read $ init $ (!! 1) $ splitOn "=" $ w !! 4
|
|
w = words s
|
|
|
|
startLocation = "AA"
|
|
totalMinutes = 30
|
|
|
|
findValve :: [Valve] -> String -> Valve
|
|
findValve valves s = fromJust $ find ((== s) . valveName) valves
|
|
|
|
calculateValvesBuildup :: [Valve] -> Int
|
|
calculateValvesBuildup = sum . map valveFlowRate . filter valveOpen
|
|
|
|
openValve :: [Valve] -> String -> [Valve]
|
|
openValve valves name = map (\v -> if valveName v == name then Valve (valveName v) (valveFlowRate v) (valveLeadsTo v) True else v) valves
|
|
|
|
pathfindValve :: [Valve] -> String -> String -> (Int, [String])
|
|
pathfindValve valves a b = fromJust $ astarSearch a (== b) (map (, 1000) . valveLeadsTo . findValve valves) (const 0)
|
|
|
|
getValveScore :: [Valve] -> String -> Int -> Valve -> (Valve, Int, [String])
|
|
getValveScore valves pos minutes valve = (valve, (totalMinutes - approxMinutes) * flowRate, path)
|
|
where
|
|
flowRate = valveFlowRate valve
|
|
approxMinutes = minutes + length path
|
|
path = snd $ pathfindValve valves pos (valveName valve)
|
|
|
|
go :: [Valve] -> Int -> String -> Int -> Int
|
|
go valves minutes pos pressure
|
|
| minutes > 30 = pressure
|
|
| null valveScores = go valves (minutes + 1) pos (pressure + calculateValvesBuildup valves)
|
|
| otherwise = maximum $ trace (map sndOf3 valveScores) $ map goTo maxValves
|
|
where
|
|
goTo :: (Valve, Int, [String]) -> Int
|
|
goTo (valve, _, path) = go newValves (minutes + 1) newPosition (pressure + calculateValvesBuildup valves)
|
|
where
|
|
newPosition = if name == pos then pos else path !! 1
|
|
newValves = if name == pos then openValve valves name else valves
|
|
name = valveName valve
|
|
|
|
maxValves = filter (\(_, score, _) -> score == maxScore) valveScores
|
|
(_, maxScore, _) = maximumBy (\a b -> sndOf3 a `compare` sndOf3 b) valveScores
|
|
valveScores = map (getValveScore valves pos minutes) $ filter (\v -> valveFlowRate v > 0 && not (valveOpen v)) valves
|
|
|
|
main = interact $
|
|
show
|
|
. (\valves -> go valves 0 startLocation 0)
|
|
. map parse
|
|
. lines
|