#568929 (Haskell) No.59 鉄道の旅

• 問題: No.59 鉄道の旅
• ユーザー: こまる
• 提出日時: 2020-10-19 13:57:10
• 言語: Haskell (9.8.2)
• 結果: CE (最新)
  AC
  (最初)
• 実行時間: -
• コード長: 3,406 bytes
• コンパイル時間: 718 ms
• コンパイル使用メモリ: 161,280 KB
• 最終ジャッジ日時: 2024-04-27 03:30:52
• 合計ジャッジ時間: 1,122 ms
• ジャッジサーバーID （参考情報）: judge1 / judge4

コンパイルメッセージ

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

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

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

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

ソースコード

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

import           Control.Monad
import           Data.Bits
import qualified Data.Array.IO                     as ArrIO
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM

main :: IO ()
main = do
  [n, k] <- map read . words <$> getLine
  bit <- newBIT 2000000
  rep n $ \_ -> do
    w <- readLn :: IO Int
    if w > 0
      then do
        temp1 <- bit -|-! 2000000
        temp2 <- bit -|-! (w - 1)
        when (temp1 - temp2 < k) $ incBIT bit w 1
      else do
        let w' = abs w
        temp3 <- bit -|-! w'
        temp4 <- bit -|-! (w' - 1)
        when (temp3 - temp4 > 0) $ incBIT bit w' (-1)
  print =<< bit -|-! 2000000

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

-------------------------------------------------------------------------------
-- Binary Indexed Tree
-------------------------------------------------------------------------------
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

type BinaryIndexedTree = ArrIO.IOUArray Int Int

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

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 = sumFromToBIT bit i (i + 1)

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

findMaxIndexLT :: BinaryIndexedTree -> Int -> IO Int
findMaxIndexLT bit w0
  | w0 <= 0   = return 0
  | otherwise = do
    n <- snd <$> ArrIO.getBounds bit
    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 findMaxIndexLT #-}