Submission #422366

---

Source Code Expand

{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NumDecimals #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExplicitForAll #-}

import Control.Applicative
import Control.Monad
import qualified Control.Monad.Primitive as Primitive
import Control.Monad.ST
import Data.Bits
import qualified Data.ByteString.Char8 as BS
import qualified Data.Foldable as Fold
import Data.List
import Data.Monoid
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import GHC.Exts (MutableByteArray#, newByteArray#, readIntArray#, writeIntArray#, Int(..))
import GHC.ST (ST(..))

main :: IO ()
main = do
  [h, _w] <- getInts
  cs <- V.replicateM h BS.getLine
  putStrLn $ case solve' cs of
    True -> "Yes"
    False -> "No"

solve' :: V.Vector BS.ByteString -> Bool
solve' grid = distanceBetween (V.map (U.map (,1)) graph) start goal < 1000
  where
    graph = gridGraph grid (/='#')
    Just start = toI <$> gridLocate grid 's'
    Just goal = toI <$> gridLocate grid 'g'
    toI (x, y) = x * w + y
    w = BS.length $ V.head grid

----------------------------------------------------------------------------
-- Grid

type Grid = V.Vector BS.ByteString

gridLocate :: Grid -> Char -> Maybe (Int, Int)
gridLocate grid ch = getFirst $ Fold.foldMap f $ V.indexed grid
  where
    f (i, row) = First $ (i,) <$> BS.findIndex (==ch) row

gridGraph :: Grid -> (Char -> Bool) -> Graph
gridGraph grid good = mkGraph (h * w) $ U.filter ok $ vert U.++ horiz
  where
    !h = V.length grid
    !w = BS.length $ V.head grid

    ok (p0, p1) = ok' p0 && ok' p1
    ok' (flip quotRem w -> (r, c)) = good $ grid V.! r `BS.index` c

    vert = U.generate ((h - 1) * w) $ \i -> let
        !i' = i + w
      in (i, i')
    horiz = U.generate (h * (w - 1)) $ \i -> let
        !(r, c) = quotRem i (w - 1)
        !p = r * w + c
        !p' = p + 1
      in (p, p')

----------------------------------------------------------------------------
-- Graph

type Graph = V.Vector (U.Vector Int)

-- | Weighted graph
type WGraph w = V.Vector (U.Vector (Int, w))

mkGraph :: Int -> U.Vector (Int, Int) -> Graph
mkGraph !nv !edges = buildGraphlike nv (2 * U.length edges) trav
  where
    trav f = U.forM_ edges $ \(a, b) -> f a b >> f b a
    {-# INLINE trav #-}

distanceBetween :: WGraph Int -> Int -> Int -> Int
distanceBetween g start goal = distancesFrom g start U.! goal

distancesFrom :: WGraph Int -> Int -> U.Vector Int
distancesFrom g start = fst $ dijkstra g start

dijkstra :: WGraph Int -> Int -> (U.Vector Int{-dist-}, U.Vector Int{-prev-})
dijkstra g start =
  runST $ gdijkstra maxBound start (nEdges+1) (V.length g) trav
  where
    nEdges = V.sum $ V.map U.length g
    trav v f = U.forM_ (g `V.unsafeIndex` v) $ \(next, cost) ->
      f next cost v
    {-# INLINE trav #-}

-- | dijkstra on an implicitly-represented graph.

-- | dijkstra on an implicitly-represented graph.
gdijkstra
  :: Int -- ^ distance assigned to unreachable nodes
  -> Int -- ^ starting vertex
  -> Int -- ^ required heap capacity
  -> Int -- ^ number of vertices
  -> (Int -> (Int{-next-} -> Int{-cost-} -> Int{-info-} -> ST s ()) -> ST s ())
    -- ^ traverse edges from the given vertex
  -> ST s (U.Vector Int, U.Vector Int)
gdijkstra !myInf !start !heapCap !sz !graph = do
  prev <- UM.new sz
  dist <- UM.replicate sz myInf
  UM.write dist start 0
  UM.write prev start (-1)
  q <- newMH heapCap
  insertMH q 0 start
  loop q prev dist
  (,) <$> U.unsafeFreeze dist <*> U.unsafeFreeze prev
  where
    loop !q !prev !dist = (deleteMH q>>=) $ Fold.mapM_ $ \(d, cur) -> do
      prevDist <- UM.read dist cur
      -- trace (printf "cur=%d d=%d prevDist=%d" cur d prevDist) $ return ()
      when (d <= prevDist) $
        graph cur $ \next cost info -> do
          nextDist <- UM.unsafeRead dist next
          let !myDist = d + cost
          -- trace (printf "cur=%d next=%d cost=%d mydist=%d" cur next cost myDist) $ return ()
          when (myDist < nextDist) $ do
            UM.unsafeWrite dist next myDist
            UM.unsafeWrite prev next info
            insertMH q myDist next
      loop q prev dist

----------------------------------------------------------------------------
-- IO

getInts :: IO [Int]
getInts = readInts <$> BS.getLine

readInts :: BS.ByteString -> [Int]
readInts = unfoldr $ \(BS.dropWhile (==' ') -> s) ->
  if s == "" then Nothing else case BS.readInt s of
    Just z -> Just z
    _ -> error $ "not an integer: " ++ show s

----------------------------------------------------------------------------
-- MHeap

data MHeap s k a = MHeap
  { mhSize :: {-# UNPACK #-} !(MBA s)
  , mhKeys :: !(UM.MVector s k)
  , mhVals :: !(UM.MVector s a)
  }

deleteMH :: (UM.Unbox k, Ord k, UM.Unbox a) => MHeap s k a -> ST s (Maybe (k, a))
deleteMH MHeap{mhSize=szV, mhVals=vals, mhKeys=keys} = do
  sz <- readIntMBA szV 0
  let !sz' = sz - 1
  if sz' < 0
    then return Nothing
    else do
      outKey <- UM.unsafeRead keys 0
      outVal <- UM.unsafeRead vals 0
      writeIntMBA szV 0 sz'
      key <- UM.unsafeRead keys sz'
      val <- UM.unsafeRead vals sz'
      loop sz' 0 key val
      return $ Just (outKey, outVal)
  where
    loop !sz !pos !(forceU -> key) !(forceU -> val)
      | lch >= sz = do
        UM.unsafeWrite keys pos key
        UM.unsafeWrite vals pos val
      | otherwise = do
        lkey <- UM.unsafeRead keys lch
        (!ch, forceU -> !ckey, forceU -> !cval) <-
          if rch >= sz
            then do
              lval <- UM.unsafeRead vals lch
              return (lch, lkey, lval)
            else do
              rkey <- UM.unsafeRead keys rch
              if lkey < rkey
                then do
                  lval <- UM.unsafeRead vals lch
                  return (lch, lkey, lval)
                else do
                  rval <- UM.unsafeRead vals rch
                  return (rch, rkey, rval)
        if key < ckey
          then do
            UM.unsafeWrite keys pos key
            UM.unsafeWrite vals pos val
          else do
            UM.unsafeWrite keys pos ckey
            UM.unsafeWrite vals pos cval
            loop sz ch key val
      where
        !lch = 2 * pos + 1
        !rch = lch + 1

insertMH :: (UM.Unbox k, Ord k, UM.Unbox a) => MHeap s k a -> k -> a -> ST s ()
insertMH MHeap{mhSize=szV, mhVals=vals, mhKeys=keys} !key !val = do
  sz <- readIntMBA szV 0
  let !sz' = sz + 1
  --trace ("insert: sz'=" ++ show sz') $ return ()
  when (sz' > UM.length vals) $ overflowErrorMH $! UM.length vals
  writeIntMBA szV 0 sz'
  loop sz
  where
    loop 0 = do
      UM.unsafeWrite keys 0 key
      UM.unsafeWrite vals 0 val
    loop pos = do
      parent <- UM.unsafeRead keys pos'
      if parent <= key
        then do
          UM.unsafeWrite keys pos key
          UM.unsafeWrite vals pos val
        else do
          UM.unsafeWrite keys pos parent
          UM.unsafeWrite vals pos =<< UM.unsafeRead vals pos'
          loop pos'
      where
        !pos' = (pos - 1) `shiftR` 1

overflowErrorMH :: Int -> ST s ()
overflowErrorMH s = fail $ "insertMH: overflow (cap=" ++ show s ++ ")"
{-# NOINLINE overflowErrorMH #-}

-- | Create an empty heap of capacity @cap@.
newMH :: (UM.Unbox k, UM.Unbox a) => Int -> ST s (MHeap s k a)
newMH cap = MHeap
  <$> do
    r <- newMBA 8
    writeIntMBA r 0 0
    return r
  <*>  UM.new cap
  <*>  UM.new cap

----------------------------------------------------------------------------
-- Util

buildGraphlike
  :: (G.Vector outv outedge)
  => Int
  -> Int
  -> (forall m. (Monad m) => (Int -> outedge -> m ()) -> m ())
  -> V.Vector (outv outedge)
buildGraphlike !nv !nEdges traverseEdges_ = runST $ do
  !counter <- UM.replicate nv 0
  let incr v _ = UM.unsafeWrite counter v . (+1) =<< UM.read counter v
  traverseEdges_ incr
  inplacePrescanl (+) 0 counter
  !mpool <- GM.new nEdges
  let
    add a oe = do
      i <- UM.unsafeRead counter a
      GM.unsafeWrite mpool i oe
      UM.unsafeWrite counter a $ i + 1
  traverseEdges_ add
  pool <- G.unsafeFreeze mpool
  V.generateM nv $ \i -> do
    begin <- if i == 0 then return 0 else UM.unsafeRead counter (i - 1)
    end <- UM.unsafeRead counter i
    return $! G.unsafeSlice begin (end - begin) pool
{-# INLINE buildGraphlike #-}

inplacePrescanl
  :: (Primitive.PrimMonad m, GM.MVector v a)
  => (a -> a -> a)
  -> a
  -> v (Primitive.PrimState m) a
  -> m ()
inplacePrescanl f x0 !mv = go 0 x0
  where
    go !i !_
      | i >= GM.length mv = return ()
    go !i !x = do
      old <- GM.unsafeRead mv i
      GM.unsafeWrite mv i x
      go (i+1) $ f x old
{-# INLINE inplacePrescanl #-}

forceU :: (U.Unbox a) => a -> a
forceU x = G.elemseq (vec x) x x
  where
    vec :: a -> U.Vector a
    vec = undefined
{-# INLINE forceU #-}

----------------------------------------------------------------------------
-- MBA

-- | Untyped byte array.
data MBA s = MBA (MutableByteArray# s)

readIntMBA :: MBA s -> Int -> ST s Int
readIntMBA (MBA mba) (I# ofs) = ST $ \s -> let
  !(# s1, val #) = readIntArray# mba ofs s
  in (# s1, I# val #)
{-# INLINE readIntMBA #-}

writeIntMBA :: MBA s -> Int -> Int -> ST s ()
writeIntMBA (MBA mba) (I# ofs) (I# val) = ST $ \s -> let
  !s1 = writeIntArray# mba ofs val s
  in (# s1, () #)
{-# INLINE writeIntMBA #-}

newMBA :: Int -> ST s (MBA s)
newMBA (I# bytes) = ST $ \s -> let
  !(# s1, mba #) = newByteArray# bytes s
  in (# s1, MBA mba #)

Submission Info

Judge Result