Skip to content

Instantly share code, notes, and snippets.

@blacktaxi blacktaxi/1930.hs
Created Oct 28, 2015

Embed
What would you like to do?
Ivan's car™
#!/usr/bin/env stack
-- stack --resolver lts-3.11 --install-ghc runghc --package array --package containers
{-# LANGUAGE ScopedTypeVariables #-}
module Main where
import qualified Data.ByteString.Char8 as B
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Data.Sequence as Seq
import Data.Sequence ((<|), (><))
import Data.Maybe
import Control.Monad (replicateM)
import Control.Applicative ((<$>))
type VertexNumber = Int
data RoadDirection = Uphill | Downhill deriving (Show)
type Edges = IM.IntMap (Seq.Seq (VertexNumber, RoadDirection))
type Graph = (VertexNumber, Edges)
type VisitedVertexes = M.Map (GearMode, VertexNumber) Cost
data GearMode = Any | Slow | Fast deriving (Show, Ord, Eq)
type Cost = Int
type SearchState = (VertexNumber, Cost, GearMode)
type Search = (Seq.Seq SearchState, VisitedVertexes)
mkGraph :: VertexNumber -> [(VertexNumber, VertexNumber)] -> Graph
mkGraph vC adj =
(vC, IM.fromListWith (><) $
concat
[ [ (a, Seq.singleton (b, Uphill)),
(b, Seq.singleton (a, Downhill)) ] | (a, b) <- adj ])
searchStep :: SearchState -> Edges -> VisitedVertexes -> (Seq.Seq SearchState, VisitedVertexes)
searchStep (v, c, g) es visited =
Seq.foldlWithIndex aux (Seq.empty, visited) nextVs
where
nextVs = fromMaybe Seq.empty (IM.lookup v es)
aux (ss, vis) _ (v', dir) =
let (cP, g') = stepCost g dir
c' = c + cP
in
case M.lookup (g', v') vis of
Nothing -> ((v', c', g') <| ss, M.insert (g', v') c' vis)
Just x | x > c' -> ((v', c', g') <| ss, M.insert (g', v') c' vis)
_ -> (ss, vis)
stepCost :: GearMode -> RoadDirection -> (Int, GearMode)
stepCost Any Uphill = (0, Slow)
stepCost Any Downhill = (0, Fast)
stepCost Slow Uphill = (0, Slow)
stepCost Slow Downhill = (1, Fast)
stepCost Fast Downhill = (0, Fast)
stepCost Fast Uphill = (1, Slow)
solve :: Graph -> (VertexNumber, VertexNumber) -> Cost
solve (_, es) (start, end) =
solve' (Seq.singleton (start, 0, Any), M.singleton (Any, start) 0)
where
solve' :: Search -> Cost
solve' (ss, visited) | Seq.null ss =
minimum $ mapMaybe (`M.lookup` visited) [(Any, end), (Slow, end), (Fast, end)]
solve' (ss, visited) =
let f (s, vis) _ x = let (s', vis') = searchStep x es vis in (s >< s', vis')
(ss', visited') = Seq.foldlWithIndex f (Seq.empty, visited) ss
in
solve' (ss', visited')
main :: IO ()
main = do
[n, m] <- readMany
(edges :: [(Int, Int)]) <- map (\[a, b] -> (a, b)) <$> replicateM m readMany
[start, end] <- readMany
print $ solve (mkGraph n edges) (start, end)
where
-- readOne = fmap (fst . fromJust . B.readInt) B.getLine
readMany = fmap (map (fst . fromJust . B.readInt) . B.words) B.getLine
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.