AtCoder Typical Contest 001

Submission #421965

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 (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状態 RE
Score得点 0
Source lengthソースコード長 9986 Byte
File nameファイル名
Exec time実行時間 ms
Memory usageメモリ使用量 -

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 0 / 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 48 ms 1260 KB
00_min_02.txt RE
00_min_03.txt RE
00_min_04.txt RE
00_min_05.txt AC 29 ms 1300 KB
00_min_06.txt RE
00_min_07.txt RE
00_min_08.txt RE
00_sample_01.txt AC 32 ms 1304 KB
00_sample_02.txt AC 30 ms 1304 KB
00_sample_03.txt AC 30 ms 1308 KB
00_sample_04.txt AC 30 ms 1312 KB
00_sample_05.txt AC 30 ms 1304 KB
01_rnd_00.txt RE
01_rnd_01.txt AC 691 ms 121368 KB
01_rnd_02.txt AC 623 ms 108124 KB
01_rnd_03.txt AC 681 ms 126488 KB
01_rnd_04.txt AC 661 ms 119460 KB
01_rnd_05.txt AC 371 ms 63124 KB
01_rnd_06.txt AC 604 ms 106776 KB
01_rnd_07.txt AC 620 ms 109212 KB
01_rnd_08.txt AC 346 ms 61460 KB
01_rnd_09.txt AC 360 ms 62484 KB
01_rnd_10.txt AC 558 ms 103188 KB
01_rnd_11.txt AC 355 ms 64540 KB
01_rnd_12.txt AC 649 ms 114836 KB
01_rnd_13.txt AC 650 ms 113952 KB
01_rnd_14.txt AC 369 ms 65176 KB
01_rnd_15.txt AC 583 ms 104728 KB
01_rnd_16.txt AC 355 ms 64536 KB
01_rnd_17.txt AC 576 ms 103768 KB
01_rnd_18.txt AC 362 ms 61464 KB
01_rnd_19.txt AC 707 ms 125336 KB
02_rndhard_00.txt AC 399 ms 65044 KB
02_rndhard_01.txt AC 373 ms 65044 KB
02_rndhard_02.txt AC 507 ms 101792 KB
02_rndhard_03.txt AC 506 ms 101788 KB
02_rndhard_04.txt AC 374 ms 65304 KB
02_rndhard_05.txt AC 374 ms 65304 KB
02_rndhard_06.txt AC 374 ms 65180 KB
02_rndhard_07.txt AC 375 ms 65172 KB
02_rndhard_08.txt AC 492 ms 101524 KB
02_rndhard_09.txt AC 489 ms 101404 KB
02_rndhard_10.txt AC 490 ms 101660 KB
02_rndhard_11.txt AC 487 ms 101652 KB
02_rndhard_12.txt AC 490 ms 101144 KB
02_rndhard_13.txt AC 488 ms 101144 KB
02_rndhard_14.txt AC 492 ms 101468 KB
02_rndhard_15.txt AC 491 ms 101400 KB
02_rndhard_16.txt AC 374 ms 65180 KB
02_rndhard_17.txt AC 373 ms 65176 KB
02_rndhard_18.txt AC 373 ms 65052 KB
02_rndhard_19.txt AC 374 ms 65052 KB
02_rndhard_20.txt AC 376 ms 65048 KB
02_rndhard_21.txt AC 374 ms 65048 KB
02_rndhard_22.txt AC 486 ms 101276 KB
02_rndhard_23.txt AC 380 ms 65048 KB
02_rndhard_24.txt AC 373 ms 65260 KB
02_rndhard_25.txt AC 373 ms 65176 KB
02_rndhard_26.txt AC 377 ms 65048 KB
02_rndhard_27.txt AC 374 ms 65044 KB
02_rndhard_28.txt AC 378 ms 64920 KB
02_rndhard_29.txt AC 378 ms 64912 KB
02_rndhard_30.txt AC 371 ms 64536 KB
02_rndhard_31.txt AC 369 ms 64536 KB
02_rndhard_32.txt AC 497 ms 101656 KB
02_rndhard_33.txt AC 496 ms 101700 KB
02_rndhard_34.txt AC 377 ms 65188 KB
02_rndhard_35.txt AC 373 ms 65168 KB
02_rndhard_36.txt AC 378 ms 65052 KB
02_rndhard_37.txt AC 371 ms 65048 KB
02_rndhard_38.txt AC 371 ms 65048 KB
02_rndhard_39.txt AC 372 ms 65048 KB
03_rndhardsmall_00.txt AC 30 ms 1304 KB
03_rndhardsmall_01.txt RE
03_rndhardsmall_02.txt AC 29 ms 1300 KB
03_rndhardsmall_03.txt AC 29 ms 1240 KB
03_rndhardsmall_04.txt AC 31 ms 1308 KB
03_rndhardsmall_05.txt AC 30 ms 1292 KB
03_rndhardsmall_06.txt AC 30 ms 1300 KB
03_rndhardsmall_07.txt AC 30 ms 1280 KB
03_rndhardsmall_08.txt AC 30 ms 1300 KB
03_rndhardsmall_09.txt AC 30 ms 1428 KB