Submission #422458

---

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 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 = distanceBetween1 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)

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 #-}

distanceBetween1 :: Graph -> Int -> Int -> Int
distanceBetween1 graph start goal = distancesFrom1 graph start U.! goal

distancesFrom1 :: Graph -> Int -> U.Vector Int
distancesFrom1 !graph !start = U.create $ do
  let !nv = V.length graph
  level <- UM.replicate nv maxBound
  q <- newMQ nv
  let
    go = (popMQ q >>=) $ Fold.mapM_ $ \v -> do
      lev <- UM.read level v
      U.forM_ (graph G.! v) $ \next -> do
        nextLev <- UM.read level next
        when (nextLev == maxBound) $ do
          UM.write level next (lev + 1)
          pushMQ q next
      go
  pushMQ q start
  UM.write level start 0
  go
  return level

----------------------------------------------------------------------------
-- 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

----------------------------------------------------------------------------
-- MQueue

data MQueue s a = MQueue
  { mqSzV :: !(MBA s) -- raedp, writep
  , mqVals :: !(UM.MVector s a)
  }

popMQ :: (UM.Unbox a) => MQueue s a -> ST s (Maybe a)
popMQ MQueue{..} = do
  readp <- readIntMBA mqSzV 0
  writep <- readIntMBA mqSzV 1
  if readp == writep
    then return Nothing
    else do
      writeIntMBA mqSzV 0 $ (readp + 1) `rem` UM.length mqVals
      Just <$> UM.unsafeRead mqVals readp
{-# INLINE popMQ #-}

pushMQ :: (UM.Unbox a) => MQueue s a -> a -> ST s ()
pushMQ MQueue{..} (forceU -> !val) = do
  readp <- readIntMBA mqSzV 0
  writep <- readIntMBA mqSzV 1
  let !writep' = rem (writep + 1) $ UM.length mqVals
  when (writep' == readp) $ fail $
    "pushMQ: overflow (cap=" ++ show (UM.length mqVals) ++ ")"
  UM.unsafeWrite mqVals writep val
  writeIntMBA mqSzV 1 writep'
{-# INLINE pushMQ #-}

newMQ :: (UM.Unbox a) => Int -> ST s (MQueue s a)
newMQ cap = MQueue
  <$> do
    r <- newMBA 16
    writeIntMBA r 0 0
    writeIntMBA r 1 0
    return r
  <*> UM.new (max 1 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