見出し画像

セグ木でデキるなら BIT でやってみる

競プロ典型 90 問 68 日目 - Paired Information(★5)

import Control.Monad
import Control.Monad.ST
import Data.Bits
import qualified Data.ByteString.Char8 as C
import Data.List
import qualified Data.Vector.Unboxed.Mutable as UM

main = readLn @Int >>= \n -> readLn >>= flip replicateM ints >>= fmt . sol n

ints = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine

fmt = mapM_ (putStrLn . maybe "Ambiguous" show)

sol n as = reverse $ runST $ do
  bt <- UM.replicate (m+1) (0,0)
  let
    dn i p
      | i>0       = UM.unsafeRead bt i >>= dn (i .&. (i-1)) . op p
      | otherwise = return p
    up i p
      | i<=m+1    = UM.unsafeModify bt (op p) i >> up (i + i .&. (-i)) p
      | otherwise = return ()
    qry l r = do
      (il,vl) <- dn (l-1) (0,0)
      (ir,vr) <- dn r (0,0)
      return (ir-il,if even l then vr-vl else vl-vr)
    go vs [0,x,_,v] = do
      qry x x >>= \(l,a) -> when (l==0) $ up x (1,if even x then v else -v)
      return vs
    go vs [_,x,y,v] = (:vs) <$!> if x==y then return (Just v) else
      qry (min x y) (pred $ max x y) >>= \(l,a) -> return (
        if l<abs (x-y) then Nothing else
          Just $ if odd (x-y)
            then a-v
            else if x>y then v+a else v-a
        )
  foldM go [] as
  where
    m = bit $ finiteBitSize n-countLeadingZeros n

op (i,v) (i',v') = (i+i',v+v')

ありがとう