#569222 (Haskell) No.833 かっこいい電車

• 問題: No.833 かっこいい電車
• ユーザー: こまる
• 提出日時: 2020-10-20 22:42:11
• 言語: Haskell (9.8.2)
• 結果: CE (最新)
  AC
  (最初)
• 実行時間: -
• コード長: 4,442 bytes
• コンパイル時間: 705 ms
• コンパイル使用メモリ: 188,732 KB
• 最終ジャッジ日時: 2024-11-14 23:52:29
• 合計ジャッジ時間: 2,072 ms
• ジャッジサーバーID （参考情報）: judge3 / judge2

コンパイルメッセージ

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

Main.hs:71:19: error: [GHC-87543]
    Ambiguous occurrence ‘.<<.’.
    It could refer to
       either ‘Data.Bits..<<.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..<<.’, defined at Main.hs:58:1.
   |
71 |   | n >= 1    = 1 .<<. (63 - (clz n))
   |                   ^^^^

Main.hs:156:37: error: [GHC-87543]
    Ambiguous occurrence ‘.>>.’.
    It could refer to
       either ‘Data.Bits..>>.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..>>.’, defined at Main.hs:62:1.
    |
156 |               then go (w - u) (step .>>. 1) (i + step) m
    |                                     ^^^^

Main.hs:157:31: error: [GHC-87543]
    Ambiguous occurrence ‘.>>.’.
    It could refer to
       either ‘Data.Bits..>>.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..>>.’, defined at Main.hs:62:1.
    |
157 |               else go w (step .>>. 1) i m
    |                               ^^^^

Main.hs:158:27: error: [GHC-87543]
    Ambiguous occurrence ‘.>>.’.
    It could refer to
       either ‘Data.Bits..>>.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..>>.’, defined at Main.hs:62:1.
    |
158 |           else go w (step .>>. 1) i m
    |                           ^^^^

ソースコード

{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase       #-}

import           Control.Arrow
import           Control.Monad
import           Control.Monad.ST
import           Control.Monad.State
import           Data.Bits
import           Data.Char
import           Data.Coerce
import qualified Data.Array.IO                     as ArrIO
import qualified Data.ByteString.Char8             as BSC8
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM
import qualified Data.Vector.Unboxed               as VU
import qualified Data.Vector.Unboxed.Mutable       as VUM

main :: IO ()
main = do
  [m, q] <- map read . words <$> getLine
  a      <- parseM m
  ba     <- newBIT m
  bb     <- newBIT m
  cs     <- VUM.replicate (m - 1) (0 :: Int)
  rep m $ \i -> do
    incBIT ba (i + 1) (a VU.! i)
    incBIT bb (i + 1) 1
  rep q $ \_ -> do
    [t, x] <- map read . words <$> getLine
    when (t == 1) $ do
      c <- VUM.unsafeRead cs (x - 1)
      when (c == 0) $ do
        incBIT bb x (-1)
        VUM.unsafeWrite cs (x - 1) 1
    when (t == 2) $ do
      c <- VUM.unsafeRead cs (x - 1)
      when (c /= 0) $ do
        incBIT bb x 1
        VUM.unsafeWrite cs (x - 1) 0
    when (t == 3) $ do
      incBIT ba x 1
    when (t == 4) $ do
      s <- bb -|-! (x - 1)
      k0 <- findMinIndexGT bb s
      k1 <- findMinIndexGT bb (s + 1)
      print =<< (-) <$> ba -|-! k1 <*> ba -|-! k0


type Parser a = BSC8.ByteString -> Maybe (a, BSC8.ByteString)
parseInt :: Parser Int
parseInt = fmap (second BSC8.tail) . BSC8.readInt
parseM :: Int -> IO (VU.Vector Int)
parseM m = VU.unfoldrN m parseInt <$> BSC8.getLine

infixl 8 .<<., .>>.

(.<<.) :: Bits b => b -> Int -> b
(.<<.) = unsafeShiftL
{-# INLINE (.<<.) #-}

(.>>.) :: Bits b => b -> Int -> b
(.>>.) = unsafeShiftR
{-# INLINE (.>>.) #-}

clz :: FiniteBits fb => fb -> Int
clz = countLeadingZeros
{-# INLINE clz #-}

ltPow2 :: Int -> Int
ltPow2 n
  | n >= 1    = 1 .<<. (63 - (clz n))
  | otherwise = 0

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

rep :: Monad m => Int -> (Int -> m ()) -> m ()
rep n = flip VFSM.mapM_ (stream 0 n)
{-# INLINE rep #-}

type BinaryIndexedTree = ArrIO.IOUArray Int Int

newBIT :: Int -> IO BinaryIndexedTree
newBIT n = ArrIO.newListArray (1, n) $ repeat 0
{-# INLINE newBIT #-}

buildBITFromVec :: VU.Vector Int -> IO BinaryIndexedTree
buildBITFromVec vec = do
  let !n = VU.length vec
  bit <- newBIT n
  rep n $ \i -> do
    incBIT bit (i + 1) (vec VU.! i)
  return bit
{-# INLINE buildBITFromVec #-}

buildBITFromList :: [Int] -> IO BinaryIndexedTree
buildBITFromList = buildBITFromVec . VU.fromList
{-# INLINE buildBITFromList #-}

infixl 9 -|-!

(-|-!) :: BinaryIndexedTree -> Int -> IO Int
bit -|-! i = iter i 0
  where
    iter :: Int -> Int -> IO Int
    iter z a
      | z < 1 = return a
      | otherwise = do
          b <- (+ a) <$> ArrIO.readArray bit z
          let j = z - (z .&. (- z))
          iter j b

sumFromToBIT :: BinaryIndexedTree -> Int -> Int -> IO Int
sumFromToBIT bit l r
  | l >= 2    = (-) <$> bit -|-! r <*> bit -|-! (l - 1)
  | otherwise = bit -|-! r

incBIT :: BinaryIndexedTree -> Int -> Int -> IO ()
incBIT bit i v = do
  (_, u) <- ArrIO.getBounds bit
  iter i u v bit
    where
      iter z key value b = when (z <= key) $ do
        ArrIO.writeArray b z . (+ value) =<< ArrIO.readArray b z
        iter (z + (z .&. (-z))) key value b

readBIT :: BinaryIndexedTree -> Int -> IO Int
readBIT bit i = (-) <$> bit -|-! i <*> bit -|-! (i - 1)

writeBIT :: BinaryIndexedTree -> Int -> Int -> IO ()
writeBIT bit i x = readBIT bit i >>= incBIT bit i . (x - )

findMinIndexGT :: BinaryIndexedTree -> Int -> IO Int
findMinIndexGT bit w0
  | w0 <= 0   = return 0
  | otherwise = do
    n    <- snd <$> ArrIO.getBounds bit
    wmax <- bit -|-! n
    if w0 > wmax
      then return (n + 1)
      else go w0 (ltPow2 n) 0 n
  where
    go !w !step !i !m
      | step == 0 = return (i + 1)
      | otherwise = do
        if i + step < m
          then do
            u <- ArrIO.readArray bit (i + step)
            if u < w
              then go (w - u) (step .>>. 1) (i + step) m
              else go w (step .>>. 1) i m
          else go w (step .>>. 1) i m
{-# INLINE findMinIndexGT #-}