2022-12-16 07:51:33 +01:00
import Control.Lens
2023-01-10 12:30:49 +01:00
import Data.Array
import Data.Int
import Data.List
2022-12-16 07:51:33 +01:00
import Data.List.Split
2023-01-10 12:30:49 +01:00
import Data.Maybe
2022-12-16 07:51:33 +01:00
2023-01-10 12:30:49 +01:00
type Valve = ( Int8 , [ Int8 ] )
2022-12-16 07:51:33 +01:00
2023-01-10 12:30:49 +01:00
type DMatrix = Array ( Int8 , Int8 ) Int8
2022-12-16 07:51:33 +01:00
main :: IO ()
main = do
2023-01-10 12:30:49 +01:00
input <- map words . lines <$> readFile " input "
let start = fromIntegral $ fromJust $ elemIndex " AA " $ map ( !! 1 ) input
let v = valveNum input
let distanceMatrix =
dijkstraAll ( buildMatrix v ) v ( pred $ fromIntegral $ length v )
print $ part1 distanceMatrix v start
print v
2022-12-16 07:51:33 +01:00
2023-01-10 12:30:49 +01:00
valveNum :: [ [ String ] ] -> [ Valve ]
valveNum input = over ( traverse . _2 . traverse ) switch $ map snd parsed
2022-12-16 07:51:33 +01:00
where
2023-01-10 12:30:49 +01:00
parsed = map parse input
valves = map fst parsed
switch x = fromIntegral $ fromJust $ elemIndex x valves
parse ( _ : name : _ : _ : rate : _ : _ : _ : _ : ts ) = ( name , ( r , nexts ) )
2022-12-16 07:51:33 +01:00
where
r = read $ init $ last $ splitOn " = " rate
nexts = map ( \ ( x : y : _ ) -> [ x , y ] ) ts
2023-01-10 12:30:49 +01:00
buildMatrix :: [ Valve ] -> DMatrix
buildMatrix valves = n
where
m =
array
( ( 0 , 0 ) , ( l , l ) )
[ ( ( x , y ) , maxBound :: Int8 ) | x <- [ 0 .. l ] , y <- [ 0 .. l ] ]
n = m // [ ( ( i , i ) , 0 ) | i <- [ 0 .. l ] ]
l = pred $ fromIntegral $ length valves
dijkstra :: DMatrix -> [ Valve ] -> [ ( Int8 , Int8 ) ] -> Int8 -> DMatrix
dijkstra m vs ( ( x , d ) : xs ) n =
dijkstra newM newVs ( xs ++ zip nexts ( repeat $ succ d ) ) n
where
newM =
m //
[ let d' = minimum [ m ! ( n , i ) , d ]
in ( ( n , i ) , d' )
| i <- nexts
]
newVs = vs & element x' .~ ( 0 , [] )
nexts = snd ( vs !! x' )
x' = fromIntegral x
dijkstra m _ [] _ = m
2022-12-16 07:51:33 +01:00
2023-01-10 12:30:49 +01:00
dijkstraAll :: DMatrix -> [ Valve ] -> Int8 -> DMatrix
dijkstraAll m vs 0 = dijkstra m vs [ ( 0 , 1 ) ] 0
dijkstraAll m vs x = dijkstraAll ( dijkstra m vs [ ( x , 1 ) ] x ) vs ( pred x )
perm30 :: DMatrix -> [ ( Int8 , Int8 ) ] -> Int -> Int8 -> [ [ ( Int8 , Int8 ) ] ]
perm30 _ [] _ _ = [ [] ]
perm30 m xxs n current
| n > 30 = [ [] ]
| otherwise =
[ y : ys | ( y , xs ) <- picks xxs , ys <- perm30 m xs ( succ n + dist y ) ( fst y ) ]
where
picks ( x : xs ) = ( x , xs ) : [ ( y , x : ys ) | ( y , ys ) <- picks xs ]
picks [] = []
dist y = fromIntegral ( m ! ( current , fst y ) )
part1 :: DMatrix -> [ Valve ] -> Int8 -> Int
part1 m valves start = maximum $ map ( calcRoute m start 30 0 ) routes
where
valves' = filter ( \ ( _ , x ) -> x /= 0 ) $ zip [ 0 .. ] $ map fst valves
routes = perm30 m valves' 0 start
calcRoute :: DMatrix -> Int8 -> Int8 -> Int -> [ ( Int8 , Int8 ) ] -> Int
calcRoute m current time acc ( node : nodes )
| newTime > 0 = calcRoute m ( fst node ) newTime newAcc nodes
| otherwise = acc
where
newTime = pred $ time - m ! ( current , fst node )
newAcc = acc + fromIntegral ( snd node ) * fromIntegral newTime
calcRoute _ _ _ acc [] = acc