Ivan's car™

1930.hs

#!/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