#580619 (Haskell) No.1285 ゴミ捨て

• 問題: No.1285 ゴミ捨て
• ユーザー: こまる
• 提出日時: 2020-11-16 17:32:38
• 言語: Haskell (9.10.1)
• 結果: AC
• 実行時間: 62 ms / 2,000 ms
• コード長: 5,698 bytes
• コンパイル時間: 14,447 ms
• コンパイル使用メモリ: 250,240 KB
• 実行使用メモリ: 12,416 KB
• 最終ジャッジ日時: 2024-11-28 21:54:18
• 合計ジャッジ時間: 15,973 ms
• ジャッジサーバーID （参考情報）: judge4 / judge3

ファイルパターン	結果
sample	AC * 2
other	AC * 22

コンパイルメッセージ

Loaded package environment from /home/judge/.ghc/x86_64-linux-9.8.2/environments/default
[1 of 2] Compiling Main             ( Main.hs, Main.o )
[2 of 2] Linking a.out

ソースコード

{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE FlexibleInstances #-}

module Main where

import           Control.Monad
import           Control.Monad.Cont
import           Control.Monad.Fix
import           Control.Monad.ST
import           Control.Monad.State
import qualified Data.ByteString.Char8             as BSC8
import           Data.Char
import           Data.Coerce
import           Data.IORef
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM

import qualified Data.Vector.Generic               as VG
import qualified Data.Vector.Generic.Mutable       as VGM
import qualified Data.Vector.Unboxed               as VU


main :: IO ()
main = do
  n <- readLn :: IO Int
  a <- quickSort <$> parseN1 n
  check <- newIORef True
  rep (n - 1) $ \i -> when (a VU.! i == a VU.! (i + 1) - 1) $ writeIORef check False
  b <- readIORef check
  if b
    then putStrLn "1"
    else putStrLn "2"

type CParser a = StateT BSC8.ByteString Maybe a
runCParser :: CParser a -> BSC8.ByteString -> Maybe (a, BSC8.ByteString)
runCParser = runStateT
{-# INLINE runCParser #-}
int :: CParser Int
int = coerce $ BSC8.readInt . BSC8.dropWhile isSpace
{-# INLINE int #-}
parseN1 :: Int -> IO (VU.Vector Int)
parseN1 n = VU.unfoldrN n (runCParser int) <$> BSC8.getContents
{-# INLINE parseN1 #-}

quickSort :: (Ord a, VG.Vector v a) => v a -> v a
quickSort = quickSortBy compare

quickSortBy :: VG.Vector v a => (a -> a -> Ordering) -> v a -> v a
quickSortBy cmp = VG.modify $ fix $ \loop vec ->
  when (VGM.length vec > 1) $ do
    pivot <- getMedian3Pivot cmp vec
    cut   <- pivotPartition cmp vec pivot
    loop (VGM.unsafeDrop cut vec)
    loop (VGM.unsafeTake cut vec)
{-# INLINE quickSortBy #-}

pivotPartition :: (VGM.MVector mv a) => (a -> a -> Ordering) -> mv s a -> a -> ST s Int
pivotPartition cmp vec pivot = fix `flip` 0 `flip` VGM.length vec $ \loop !l !r -> do
  !l' <- flip fix l $ \loopL !i -> do
    x   <- VGM.unsafeRead vec i
    case cmp x pivot of
      LT -> loopL (i + 1)
      _  -> return i
  !r' <- flip fix (r - 1) $ \loopR !i -> do
    x <- VGM.unsafeRead vec i
    case cmp pivot x of
      LT -> loopR (i - 1)
      _  -> return i
  if l' < r'
    then do
      VGM.unsafeSwap vec l' r'
      loop (l' + 1) r'
    else return l'
{-# INLINE pivotPartition #-}

getMedian3Pivot :: (VGM.MVector mv a) => (a -> a -> Ordering) -> mv s a -> ST s a
getMedian3Pivot cmp vec = median cmp <$> VGM.unsafeRead vec 0 <*> VGM.unsafeRead vec (VGM.length vec `quot` 2) <*> VGM.unsafeRead vec (VGM.length vec - 1)
{-# INLINE getMedian3Pivot #-}

median :: (a -> a -> Ordering) -> a -> a -> a -> a
median cmp x y z = case cmp x y of
  LT -> case cmp y z of
    LT -> y
    _  -> case cmp x z of
      LT -> z
      _  -> x
  _  -> case cmp x z of
    LT -> x
    _  -> case cmp y z of
      LT -> z
      _  -> y
{-# INLINE median #-}

-- | l -> x -> r, +d
stream :: Monad m => Int -> Int -> Int -> VFSM.Stream m Int
stream !l !r !d = VFSM.Stream step l
  where
    step x
      | x <= r    = return $ VFSM.Yield x (x + d)
      | otherwise = return VFSM.Done
    {-# INLINE [0] step #-}
{-# INLINE [1] stream #-}

-- | 0 <= x < n, interval = 1
rep :: Monad m => Int -> (Int -> m ()) -> m ()
rep n = flip VFSM.mapM_ (stream 0 (n - 1) 1)
{-# INLINE rep #-}

-- | 0 <= x <= n, interval = 1
rep' :: Monad m => Int -> (Int -> m ()) -> m ()
rep' n = flip VFSM.mapM_ (stream 0 n 1)
{-# INLINE rep' #-}

-- | 1 <= x < n, interval = 1
rep1 :: Monad m => Int -> (Int -> m ()) -> m ()
rep1 n = flip VFSM.mapM_ (stream 1 (n - 1) 1)
{-# INLINE rep1 #-}

-- | 1 <= x <= n, interval = 1
rep1' :: Monad m => Int -> (Int -> m ()) -> m ()
rep1' n = flip VFSM.mapM_ (stream 1 n 1)
{-# INLINE rep1' #-}

-- | l <= x <= r, interval = d
for :: Monad m => Int -> Int -> Int -> (Int -> m ()) -> m ()
for l r d = flip VFSM.mapM_ (stream l r d)
{-# INLINE for #-}

-- | r -> x -> l, -d
streamR :: Monad m => Int -> Int -> Int -> VFSM.Stream m Int
streamR !r !l !d = VFSM.Stream step r
  where
    step x
      | x >= l    = return $ VFSM.Yield x (x - d)
      | otherwise = return VFSM.Done
    {-# INLINE [0] step #-}
{-# INLINE [1] streamR #-}

-- | n > x >= 0, interval = -1
rev :: Monad m => Int -> (Int -> m ()) -> m ()
rev n = flip VFSM.mapM_ (streamR (n - 1) 0 1)
{-# INLINE rev #-}

-- | n >= x >= 0, interval = -1
rev' :: Monad m => Int -> (Int -> m ()) -> m ()
rev' n = flip VFSM.mapM_ (streamR n 0 1)
{-# INLINE rev' #-}

-- | n > x >= 1, interval = -1
rev1 :: Monad m => Int -> (Int -> m ()) -> m ()
rev1 n = flip VFSM.mapM_ (streamR (n - 1) 1 1)
{-# INLINE rev1 #-}

-- | n >= x >= 1, interval = -1
rev1' :: Monad m => Int -> (Int -> m ()) -> m ()
rev1' n = flip VFSM.mapM_ (streamR n 1 1)
{-# INLINE rev1' #-}

-- | r >= x >= l, interval = -d
forR :: Monad m => Int -> Int -> Int -> (Int -> m ()) -> m ()
forR r l d = flip VFSM.mapM_ (streamR r l d)
{-# INLINE forR #-}

-- | for (int i = l; f(i, p) <= r ; g(i, d))
streamG :: Monad m => Int -> Int -> (Int -> Int -> Int) -> Int -> (Int -> Int -> Int) -> Int -> VFSM.Stream m Int
streamG l r f p g d = VFSM.Stream step l
  where
    step x
      | f x p <= r = return $ VFSM.Yield x (g x d)
      | otherwise  = return VFSM.Done
    {-# INLINE [0] step #-}
{-# INLINE [1] streamG #-}

forG :: Monad m => Int -> Int -> (Int -> Int -> Int) -> Int -> (Int -> Int -> Int) -> Int -> (Int -> m ()) -> m ()
forG l r f p g d = flip VFSM.mapM_ (streamG l r f p g d)
{-# INLINE forG #-}

withBreakIO :: ((r -> ContT r IO b) -> ContT r IO r) -> IO r
withBreakIO = flip runContT pure . callCC
{-# INLINE withBreakIO #-}

withBreakST :: ((r -> ContT r (ST s) b) -> ContT r (ST s) r) -> (ST s) r
withBreakST = flip runContT pure . callCC
{-# INLINE withBreakST #-}