Submission #52906621

---

Source Code Expand

import qualified Data.ByteString.Char8 as BS
import Data.Char
import Data.List

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

import Control.Monad
import Control.Monad.ST
import qualified Data.Vector.Unboxed.Mutable as MUV

main :: IO ()
main = do
  [n] <- bsGetLnInts
  as <- bsGetLnInts
  let ans = abc351f n as
  print ans

bsGetLnInts :: IO [Int]
bsGetLnInts = unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine

abc351f :: Int -> [Int] -> Int
abc351f n as = ans
  where
-- Aiの大きい順に背番号
--    a2i = IM.fromList $ zip (IS.toDescList $ IS.fromList as) [0..]
    a2i = IM.fromSet (\k -> maybe 0 (succ . snd) $ IM.lookupGT k a2i) $ IS.fromList as
-- Aiの異なる要素の種類数
    m = succ $ snd $ IM.findMin a2i
-- 総和と個数の両方を足す演算
    op (a,b) (c,d) = (a+c, b+d)
    ans = runST $ do
      st <- makeSegTreeN op (0,0) m :: ST s (SegmentTree (ST s) (Int,Int))
      recur st as
-- リストの後ろから計算する
    recur _st [] = return 0
    recur st (a:as) = do
      acc <- recur st as
      let i = a2i IM.! a
      (v,k) <- querySegTree st 0 i
      modifySegTree st i (op (a, 1))
      return $! acc + v - k * a

data SegmentTree m a = SegmentTree Int (a->a->a) a (MUV.MVector (MUV.PrimState m) a)

{- makeSegTreeN len op u 要素数を指定して単位元で満たされたセグメント木を作る
要素数len 添え字0始まり
op モノイド演算
u opの単位元
-}
makeSegTreeN :: (MUV.Unbox a, MUV.PrimMonad m) => (a->a->a) -> a -> Int -> m (SegmentTree m a)
makeSegTreeN op u len = SegmentTree w op u <$> MUV.replicate (w + pred w) u
  where w = until (len <=) (2 *) 1

{- querySegTree st a b [a, b)の区間をopした結果を求める -}
querySegTree :: (MUV.Unbox a, MUV.PrimMonad m) => SegmentTree m a -> Int -> Int -> m a
querySegTree (SegmentTree w op u arr) a b = loop arr 0 w 0
  where
    loop arr p w i
      | q <= a || b <= p = return u
      | a <= p && q <= b = MUV.read arr i
      | otherwise = do
        l <- loop arr  p       w2 (i + succ i)
        r <- loop arr (p + w2) w2 (2 * succ i)
        return (op l r)
      where
        q = p + w
        w2 = div w 2

{- modifySegTree st i f 要素iにfを適用する -}
modifySegTree :: (MUV.Unbox a, MUV.PrimMonad m) => SegmentTree m a -> Int -> (a -> a) -> m ()
modifySegTree (SegmentTree w op _ arr) j f =
  do
    MUV.modify arr f i0
    forM_ upwards (\i -> do
      l <- MUV.read arr (i + succ i)
      r <- MUV.read arr (2 * succ i)
      MUV.write arr i (op l r)
      )
  where
    i0 = j + pred w
    upwards = tail $ takeWhile (0 <=) $ iterate (flip div 2 . pred) i0

Submission Info

Compile Error

app/Main.hs:23:9: warning: [-Wunused-matches]
    Defined but not used: ‘n’
   |
23 | abc351f n as = ans
   |         ^

app/Main.hs:37:17: warning: [-Wname-shadowing]
    This binding for ‘as’ shadows the existing binding
      bound at app/Main.hs:23:11
   |
37 |     recur st (a:as) = do
   |                 ^^

app/Main.hs:59:10: warning: [-Wname-shadowing]
    This binding for ‘arr’ shadows the existing binding
      bound at app/Main.hs:57:34
   |
59 |     loop arr p w i
   |          ^^^

app/Main.hs:59:16: warning: [-Wname-shadowing]
    This binding for ‘w’ shadows the existing binding
      bound at app/Main.hs:57:27
   |
59 |     loop arr p w i
   |                ^

Judge Result