{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TupleSections #-} -- module Yukicoder.No1268 where import Control.Monad (when) import Data.Bits (Bits (unsafeShiftL, unsafeShiftR, (.&.), (.|.))) import Data.Char (isSpace) import Data.Coerce (coerce) import qualified Data.Foldable as F import Data.IORef (modifyIORef', newIORef, readIORef) import Data.Maybe (fromJust) import Data.Word (Word32, Word64, Word8) import Unsafe.Coerce (unsafeCoerce) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BSC8 import Control.Monad.State (StateT (..), modify', when) 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 n <- readLn :: IO Int as <- seqInput n let ax = radixSortNonNegative as dp <- VUM.replicate (n + 1) 0 :: IO (VUM.IOVector Int) rev' n $ \i -> do when (i + 1 < n && (2 + ax VU.! i) == ax VU.! (i + 1)) $ do dpi1 <- VUM.unsafeRead dp (i + 1) VUM.unsafeWrite dp i (dpi1 + 1) when (i + 2 < n && (2 + ax VU.! i) == ax VU.! (i + 2)) $ do dpi2 <- VUM.unsafeRead dp (i + 2) VUM.unsafeWrite dp i (dpi2 + 1) ans <- newIORef n rep (n - 1) $ \i -> do when (ax VU.! i + 1 == ax VU.! (i + 1)) $ do dpi1 <- VUM.unsafeRead dp (i + 1) modifyIORef' ans (+ (dpi1 + 1)) readIORef ans >>= print 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, rep', rep1, rep1' :: Monad m => Int -> (Int -> m ()) -> m () rep n = flip VFSM.mapM_ (stream 0 n) {-# INLINE rep #-} rep' n = flip VFSM.mapM_ (stream 0 (n + 1)) {-# INLINE rep' #-} rep1 n = flip VFSM.mapM_ (stream 1 n) {-# INLINE rep1 #-} rep1' n = flip VFSM.mapM_ (stream 1 (n + 1)) {-# INLINE rep1' #-} streamR :: Monad m => Int -> Int -> VFSM.Stream m Int streamR !l !r = VFSM.Stream step (r - 1) where step x | x >= l = return $ VFSM.Yield x (x - 1) | otherwise = return VFSM.Done {-# INLINE [0] step #-} {-# INLINE [1] streamR #-} rev, rev', rev1, rev1' :: Monad m => Int -> (Int -> m ()) -> m () rev !n = flip VFSM.mapM_ (streamR 0 n) {-# INLINE rev #-} rev' !n = flip VFSM.mapM_ (streamR 0 (n + 1)) {-# INLINE rev' #-} rev1 !n = flip VFSM.mapM_ (streamR 1 n) {-# INLINE rev1 #-} rev1' !n = flip VFSM.mapM_ (streamR 1 (n + 1)) {-# INLINE rev1' #-} infixl 8 .<<., .>>. (.<<.) :: Bits b => b -> Int -> b (.<<.) = unsafeShiftL {-# INLINE (.<<.) #-} (.>>.) :: Bits b => b -> Int -> b (.>>.) = unsafeShiftR {-# INLINE (.>>.) #-} class Word64Encode a where encode64 :: a -> Word64 decode64 :: Word64 -> a encodeNonNegative64 :: a -> Word64 encodeNonNegative64 = encode64 decodeNonNegative64 :: Word64 -> a decodeNonNegative64 = decode64 instance Word64Encode Int where encode64 x = unsafeCoerce $ x + 0x3fffffffffffffff decode64 x = unsafeCoerce x - 0x3fffffffffffffff encodeNonNegative64 = unsafeCoerce decodeNonNegative64 = unsafeCoerce instance Word64Encode (Int, Int) where encode64 (x, y) = unsafeCoerce $ (x + 0x3fffffff) .<<. 31 .|. (y + 0x3fffffff) decode64 xy = unsafeCoerce (x, y) where !x = xy .>>. 31 - 0x3fffffff !y = (xy .&. 0x7fffffff) - 0x3fffffff encodeNonNegative64 (x, y) = unsafeCoerce $ x .<<. 31 .|. y decodeNonNegative64 xy = unsafeCoerce (x, y) where !x = xy .>>. 31 !y = xy .&. 0x7fffffff instance Word64Encode (Int, Int, Int) where encode64 (x, y, z) = unsafeCoerce $ ((x + 0xfffff) .<<. 21 .|. (y + 0xfffff)) .<<. 21 .|. (z + 0xfffff) decode64 xyz = unsafeCoerce (x, y, z) where !x = xyz .>>. 42 - 0xfffff !y = (xyz .>>. 21 .&. 0x1fffff) - 0xfffff !z = xyz .&. 0x1fffff - 0xfffff encodeNonNegative64 (x, y, z) = unsafeCoerce $ (x .<<. 21 .|. y) .<<. 21 .|. z decodeNonNegative64 xyz = unsafeCoerce (x, y, z) where !x = xyz .>>. 42 !y = xyz .>>. 21 .&. 0x1fffff !z = xyz .&. 0x1fffff radixSortInt :: VU.Vector Int -> VU.Vector Int radixSortInt = unsafeCoerce . radixSort64 . unsafeCoerce radixSort64 :: VU.Vector Word64 -> VU.Vector Word64 radixSort64 vword = F.foldl' step vword [0,16,32,48] where mask k x = fromIntegral $ x .>>. k .&. 0xffff step v k = VU.create $ do pref <- VU.unsafeThaw . VU.prescanl' (+) 0 . VU.unsafeAccumulate (+) (VU.replicate 0x10000 0) $ VU.map ((, 1) . mask k) v res <- VUM.unsafeNew $ VU.length v VU.forM_ v $ \x -> do let !masked = mask k x i <- VUM.unsafeRead pref masked VUM.unsafeWrite pref masked $ i + 1 VUM.unsafeWrite res i x return res {-# INLINE radixSort64 #-} radixSort :: (VU.Unbox a, Word64Encode a) => VU.Vector a -> VU.Vector a radixSort = VU.map decode64 . radixSort64 . VU.map encode64 {-# INLINE radixSort #-} radixSortNonNegative :: (VU.Unbox a, Word64Encode a) => VU.Vector a -> VU.Vector a radixSortNonNegative = VU.map decodeNonNegative64 . radixSort64 . VU.map encodeNonNegative64 {-# INLINE radixSortNonNegative #-} radixSort32 :: VU.Vector Word32 -> VU.Vector Word32 radixSort32 vec = F.foldl' step vec [0, 16] where mask k x = fromIntegral $ x .>>. k .&. 0xffff step v k = VU.create $ do pref <- VU.unsafeThaw . VU.prescanl' (+) 0 . VU.unsafeAccumulate (+) (VU.replicate 0x10000 0) $ VU.map ((, 1) . mask k) v res <- VUM.unsafeNew $ VU.length v VU.forM_ v $ \x -> do let !masked = mask k x i <- VUM.unsafeRead pref masked VUM.unsafeWrite pref masked $ i + 1 VUM.unsafeWrite res i x return res {-# INLINE radixSort32 #-} compress :: VU.Vector Int -> VU.Vector Int compress vec = VU.create $ do mvec <- VUM.unsafeNew (VU.length vec) VU.mapM_ (\(i, x) -> VUM.unsafeWrite mvec (x .&. 0xffffffff) i) . VU.postscanl' (\(!i, !x) y -> if x .>>. 32 == y .>>. 32 then (i, y) else (i + 1, y) ) (-1, -1) . radixSortInt $ VU.imap (\i x -> x .<<. 32 .|. i) vec return mvec {-# INLINE compress #-} type Parser a = StateT BSC8.ByteString Maybe a runParser :: Parser a -> BSC8.ByteString -> Maybe (a, BSC8.ByteString) runParser = runStateT {-# INLINE runParser #-} int :: Parser Int int = coerce $ BSC8.readInt . BSC8.dropWhile isSpace {-# INLINE int #-} int1 :: Parser Int int1 = fmap (subtract 1) int {-# INLINE int1 #-} char :: Parser Char char = coerce BSC8.uncons {-# INLINE char #-} byte :: Parser Word8 byte = coerce BS.uncons {-# INLINE byte #-} skipSpaces :: Parser () skipSpaces = modify' (BSC8.dropWhile isSpace) {-# INLINE skipSpaces #-} seqInput :: Int -> IO (VU.Vector Int) seqInput n = VU.unfoldrN n (runParser int) <$> BSC8.getLine {-# INLINE seqInput #-} parseN1 :: Int -> IO (VU.Vector Int) parseN1 n = VU.unfoldrN n (runParser int) <$> BSC8.getContents {-# INLINE parseN1 #-} parseN2 :: Int -> IO (VU.Vector (Int, Int)) parseN2 n = VU.unfoldrN n (runParser $ (,) <$> int <*> int) <$> BSC8.getContents {-# INLINE parseN2 #-} parseN3 :: Int -> IO (VU.Vector (Int, Int, Int)) parseN3 n = VU.unfoldrN n (runParser $ (,,) <$> int <*> int <*> int) <$> BSC8.getContents {-# INLINE parseN3 #-} parseN4 :: Int -> IO (VU.Vector (Int, Int, Int, Int)) parseN4 n = VU.unfoldrN n (runParser $ (,,,) <$> int <*> int <*> int <*> int) <$> BSC8.getContents {-# INLINE parseN4 #-} parseN5 :: Int -> IO (VU.Vector (Int, Int, Int, Int, Int)) parseN5 n = VU.unfoldrN n (runParser $ (,,,,) <$> int <*> int <*> int <*> int <*> int) <$> BSC8.getContents {-# INLINE parseN5 #-} parseANBN :: Int -> IO (VU.Vector Int, VU.Vector Int) parseANBN n = VU.unzip . VU.unfoldrN n (runParser $ (,) <$> int <*> int) <$> BSC8.getContents {-# INLINE parseANBN #-} parseANBNCN :: Int -> IO (VU.Vector Int, VU.Vector Int, VU.Vector Int) parseANBNCN n = VU.unzip3 . VU.unfoldrN n (runParser $ (,,) <$> int <*> int <*> int) <$> BSC8.getContents {-# INLINE parseANBNCN #-} readInt :: BSC8.ByteString -> Int readInt = fst . fromJust . BSC8.readInt {-# INLINE readInt #-} getInt :: IO Int getInt = readInt <$> BSC8.getLine {-# INLINE getInt #-} readIntList :: BSC8.ByteString -> [Int] readIntList = map readInt . BSC8.words {-# INLINE readIntList #-} getIntList :: IO [Int] getIntList = readIntList <$> BSC8.getLine {-# INLINE getIntList #-} readInteger :: BSC8.ByteString -> Integer readInteger = fst . fromJust . BSC8.readInteger {-# INLINE readInteger #-} getInteger :: IO Integer getInteger = readInteger <$> BSC8.getLine {-# INLINE getInteger #-} readIntegerList :: BSC8.ByteString -> [Integer] readIntegerList = map readInteger . BSC8.words {-# INLINE readIntegerList #-} getIntegerList :: IO [Integer] getIntegerList = readIntegerList <$> BSC8.getLine {-# INLINE getIntegerList #-}