AtCoder Typical Contest 001

Submission #422366

Source codeソースコード

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

Task問題 A - 深さ優先探索
User nameユーザ名 mkotha
Created time投稿日時
Language言語 Haskell (Haskell Platform 2014.2.0.0)
Status状態 AC
Score得点 100
Source lengthソースコード長 9990 Byte
File nameファイル名
Exec time実行時間 698 ms
Memory usageメモリ使用量 126288 KB

Test case

Set

Set name Score得点 / Max score Cases
Sample - 00_sample_01.txt,00_sample_02.txt,00_sample_03.txt,00_sample_04.txt,00_sample_05.txt
All 100 / 100 00_min_01.txt,00_min_02.txt,00_min_03.txt,00_min_04.txt,00_min_05.txt,00_min_06.txt,00_min_07.txt,00_min_08.txt,00_sample_01.txt,00_sample_02.txt,00_sample_03.txt,00_sample_04.txt,00_sample_05.txt,01_rnd_00.txt,01_rnd_01.txt,01_rnd_02.txt,01_rnd_03.txt,01_rnd_04.txt,01_rnd_05.txt,01_rnd_06.txt,01_rnd_07.txt,01_rnd_08.txt,01_rnd_09.txt,01_rnd_10.txt,01_rnd_11.txt,01_rnd_12.txt,01_rnd_13.txt,01_rnd_14.txt,01_rnd_15.txt,01_rnd_16.txt,01_rnd_17.txt,01_rnd_18.txt,01_rnd_19.txt,02_rndhard_00.txt,02_rndhard_01.txt,02_rndhard_02.txt,02_rndhard_03.txt,02_rndhard_04.txt,02_rndhard_05.txt,02_rndhard_06.txt,02_rndhard_07.txt,02_rndhard_08.txt,02_rndhard_09.txt,02_rndhard_10.txt,02_rndhard_11.txt,02_rndhard_12.txt,02_rndhard_13.txt,02_rndhard_14.txt,02_rndhard_15.txt,02_rndhard_16.txt,02_rndhard_17.txt,02_rndhard_18.txt,02_rndhard_19.txt,02_rndhard_20.txt,02_rndhard_21.txt,02_rndhard_22.txt,02_rndhard_23.txt,02_rndhard_24.txt,02_rndhard_25.txt,02_rndhard_26.txt,02_rndhard_27.txt,02_rndhard_28.txt,02_rndhard_29.txt,02_rndhard_30.txt,02_rndhard_31.txt,02_rndhard_32.txt,02_rndhard_33.txt,02_rndhard_34.txt,02_rndhard_35.txt,02_rndhard_36.txt,02_rndhard_37.txt,02_rndhard_38.txt,02_rndhard_39.txt,03_rndhardsmall_00.txt,03_rndhardsmall_01.txt,03_rndhardsmall_02.txt,03_rndhardsmall_03.txt,03_rndhardsmall_04.txt,03_rndhardsmall_05.txt,03_rndhardsmall_06.txt,03_rndhardsmall_07.txt,03_rndhardsmall_08.txt,03_rndhardsmall_09.txt

Test case

Case name Status状態 Exec time実行時間 Memory usageメモリ使用量
00_min_01.txt AC 29 ms 1052 KB
00_min_02.txt AC 27 ms 1044 KB
00_min_03.txt AC 28 ms 980 KB
00_min_04.txt AC 29 ms 1044 KB
00_min_05.txt AC 31 ms 1052 KB
00_min_06.txt AC 31 ms 1008 KB
00_min_07.txt AC 28 ms 980 KB
00_min_08.txt AC 27 ms 976 KB
00_sample_01.txt AC 29 ms 1048 KB
00_sample_02.txt AC 29 ms 1056 KB
00_sample_03.txt AC 29 ms 1060 KB
00_sample_04.txt AC 29 ms 1176 KB
00_sample_05.txt AC 29 ms 980 KB
01_rnd_00.txt AC 358 ms 64280 KB
01_rnd_01.txt AC 698 ms 121172 KB
01_rnd_02.txt AC 612 ms 107920 KB
01_rnd_03.txt AC 693 ms 126288 KB
01_rnd_04.txt AC 667 ms 119256 KB
01_rnd_05.txt AC 373 ms 62928 KB
01_rnd_06.txt AC 611 ms 106580 KB
01_rnd_07.txt AC 629 ms 109008 KB
01_rnd_08.txt AC 353 ms 61268 KB
01_rnd_09.txt AC 364 ms 62160 KB
01_rnd_10.txt AC 565 ms 102992 KB
01_rnd_11.txt AC 358 ms 64336 KB
01_rnd_12.txt AC 657 ms 114660 KB
01_rnd_13.txt AC 654 ms 113616 KB
01_rnd_14.txt AC 378 ms 64980 KB
01_rnd_15.txt AC 589 ms 104408 KB
01_rnd_16.txt AC 365 ms 64340 KB
01_rnd_17.txt AC 599 ms 103508 KB
01_rnd_18.txt AC 355 ms 61272 KB
01_rnd_19.txt AC 674 ms 125076 KB
02_rndhard_00.txt AC 381 ms 64856 KB
02_rndhard_01.txt AC 378 ms 64856 KB
02_rndhard_02.txt AC 520 ms 101592 KB
02_rndhard_03.txt AC 510 ms 101584 KB
02_rndhard_04.txt AC 380 ms 64980 KB
02_rndhard_05.txt AC 380 ms 64980 KB
02_rndhard_06.txt AC 379 ms 64972 KB
02_rndhard_07.txt AC 386 ms 64984 KB
02_rndhard_08.txt AC 499 ms 101332 KB
02_rndhard_09.txt AC 495 ms 101204 KB
02_rndhard_10.txt AC 496 ms 101464 KB
02_rndhard_11.txt AC 498 ms 101460 KB
02_rndhard_12.txt AC 497 ms 100952 KB
02_rndhard_13.txt AC 496 ms 100952 KB
02_rndhard_14.txt AC 496 ms 101196 KB
02_rndhard_15.txt AC 497 ms 101208 KB
02_rndhard_16.txt AC 376 ms 64980 KB
02_rndhard_17.txt AC 375 ms 64980 KB
02_rndhard_18.txt AC 377 ms 64856 KB
02_rndhard_19.txt AC 375 ms 64852 KB
02_rndhard_20.txt AC 376 ms 64852 KB
02_rndhard_21.txt AC 374 ms 64852 KB
02_rndhard_22.txt AC 496 ms 101076 KB
02_rndhard_23.txt AC 378 ms 64848 KB
02_rndhard_24.txt AC 380 ms 65052 KB
02_rndhard_25.txt AC 379 ms 64976 KB
02_rndhard_26.txt AC 381 ms 64856 KB
02_rndhard_27.txt AC 377 ms 64852 KB
02_rndhard_28.txt AC 374 ms 64724 KB
02_rndhard_29.txt AC 374 ms 64720 KB
02_rndhard_30.txt AC 375 ms 64344 KB
02_rndhard_31.txt AC 375 ms 64340 KB
02_rndhard_32.txt AC 505 ms 101464 KB
02_rndhard_33.txt AC 501 ms 101468 KB
02_rndhard_34.txt AC 379 ms 64852 KB
02_rndhard_35.txt AC 378 ms 64852 KB
02_rndhard_36.txt AC 377 ms 64848 KB
02_rndhard_37.txt AC 378 ms 64848 KB
02_rndhard_38.txt AC 379 ms 64792 KB
02_rndhard_39.txt AC 376 ms 64856 KB
03_rndhardsmall_00.txt AC 28 ms 1064 KB
03_rndhardsmall_01.txt AC 29 ms 1076 KB
03_rndhardsmall_02.txt AC 29 ms 1104 KB
03_rndhardsmall_03.txt AC 29 ms 1076 KB
03_rndhardsmall_04.txt AC 29 ms 1068 KB
03_rndhardsmall_05.txt AC 29 ms 1044 KB
03_rndhardsmall_06.txt AC 29 ms 1044 KB
03_rndhardsmall_07.txt AC 29 ms 1060 KB
03_rndhardsmall_08.txt AC 29 ms 1072 KB
03_rndhardsmall_09.txt AC 29 ms 1044 KB