結果
問題 | No.206 数の積集合を求めるクエリ |
ユーザー | かりあげクン |
提出日時 | 2020-09-02 21:51:36 |
言語 | Haskell (9.8.2) |
結果 |
CE
(最新)
AC
(最初)
|
実行時間 | - |
コード長 | 16,362 bytes |
コンパイル時間 | 272 ms |
コンパイル使用メモリ | 170,112 KB |
最終ジャッジ日時 | 2024-11-14 23:48:42 |
合計ジャッジ時間 | 787 ms |
ジャッジサーバーID (参考情報) |
judge2 / judge3 |
(要ログイン)
コンパイルエラー時のメッセージ・ソースコードは、提出者また管理者しか表示できないようにしております。(リジャッジ後のコンパイルエラーは公開されます)
ただし、clay言語の場合は開発者のデバッグのため、公開されます。
ただし、clay言語の場合は開発者のデバッグのため、公開されます。
コンパイルメッセージ
Loaded package environment from /home/judge/.ghc/x86_64-linux-9.8.2/environments/default Main.hs:9:28: warning: [GHC-53692] [-Wdeprecated-flags] -XTypeInType is deprecated: use -XDataKinds and -XPolyKinds instead | 9 | {-# LANGUAGE TypeFamilies, TypeInType, UnboxedTuples, ViewPatterns #-} | ^^^^^^^^^^ [1 of 2] Compiling Main ( Main.hs, Main.o ) Main.hs:44:1: error: [GHC-61948] Could not find module ‘Data.ByteString.Lazy.Builder’. Perhaps you meant Data.ByteString.Builder (from bytestring-0.12.1.0) Data.ByteString.Lazy.Char8 (from bytestring-0.12.1.0) Data.ByteString.Lazy.ReadInt Use -v to see a list of the files searched for. | 44 | import qualified Data.ByteString.Lazy.Builder as BSLB | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Main.hs:48:1: error: [GHC-87110] Could not load module ‘Data.Graph’. It is a member of the hidden package ‘containers-0.6.8’. Use -v to see a list of the files searched for. | 48 | import qualified Data.Graph as Graph | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Main.hs:49:1: error: [GHC-87110] Could not load module ‘Data.IntMap’. It is a member of the hidden package ‘containers-0.6.8’. Use -v to see a list of the files searched for. | 49 | import Data.IntMap (IntMap) | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Main.hs:50:1: error: [GHC-87110] Could not load module ‘Data.IntMap’. It is a member of the hidden package ‘containers-0.6.8’. Use -v to see a list of the files searched for. | 50 | import qualified Data.IntMap as IntMap | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Main.hs:51:1: error: [GHC-87110] Could not load module ‘Data.IntSet’. It is a member of the hidden package ‘containers-0.6.8’. Use -v to see a list of the files searc
ソースコード
{-# LANGUAGE DeriveFunctor #-} {-# LANGUAGE BangPatterns, BinaryLiterals, CPP, DerivingStrategies #-} {-# LANGUAGE DerivingVia, FlexibleContexts, FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving, KindSignatures, LambdaCase #-} {-# LANGUAGE MagicHash, MultiParamTypeClasses, MultiWayIf #-} {-# LANGUAGE NumericUnderscores, OverloadedStrings, PatternSynonyms #-} {-# LANGUAGE RankNTypes, RecordWildCards, ScopedTypeVariables #-} {-# LANGUAGE StandaloneDeriving, TupleSections, TypeApplications #-} {-# LANGUAGE TypeFamilies, TypeInType, UnboxedTuples, ViewPatterns #-} {- base -} import Control.Applicative import qualified Control.Arrow as Arrow import Control.Monad import Control.Monad.ST import Data.Bits import Data.Bool import Data.Complex import qualified Data.Char as Char import qualified Data.Foldable as Foldable import Data.Function import qualified Data.List as List import Data.Maybe import Data.Monoid import Data.Ratio import Data.Ord import Data.Semigroup import qualified Data.Word as Word import Foreign hiding (void) import GHC.Exts import Unsafe.Coerce {- array -} import qualified Data.Array.IO as ArrIO import qualified Data.Array.MArray as ArrMA import qualified Data.Array.ST as ArrST import qualified Data.Array.Storable as ArrStore import qualified Data.Array.Unboxed as ArrU {- bytestring -} import qualified Data.ByteString as BS import qualified Data.ByteString.Builder as BSB import qualified Data.ByteString.Builder.Extra as BSBE import qualified Data.ByteString.Char8 as BSC8 import qualified Data.ByteString.Lazy as BSL import qualified Data.ByteString.Lazy.Builder as BSLB import qualified Data.ByteString.Lazy.Char8 as BSLC8 import qualified Data.ByteString.Unsafe as BSU {- containers -} import qualified Data.Graph as Graph import Data.IntMap (IntMap) import qualified Data.IntMap as IntMap import Data.IntSet (IntSet) import qualified Data.IntSet as IntSet import qualified Data.Sequence as Seq import qualified Data.Tree as Tree {- integer-gmp -} import GHC.Integer.GMP.Internals {- vector -} import qualified Data.Vector as V import qualified Data.Vector.Generic as VG import qualified Data.Vector.Generic.Mutable as VGM import qualified Data.Vector.Mutable as VM import qualified Data.Vector.Primitive as VP import qualified Data.Vector.Primitive.Mutable as VPM import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM main :: IO () main = do (l,m,n) <- parse3 a <- parseM l b <- parseM m q <- parse1 ax <- VUM.new 262144 forM_ [0..l-1] $ \i -> do VUM.unsafeWrite ax (a VU.! i - 1) (1 :: Int) as <- VU.unsafeFreeze ax bx <- VUM.new 262144 forM_ [0..m-1] $ \i -> do VUM.unsafeWrite bx (n - b VU.! i) (1 :: Int) bs <- VU.unsafeFreeze bx BSC8.putStr . BSC8.unlines . map (BSC8.pack . show) . VU.toList . VU.drop (n - 1) . VU.take (n + q - 1) $ multiply as bs type Parser a = BSC8.ByteString -> Maybe (a, BSC8.ByteString) parseInt :: Parser Int parseInt = fmap (Arrow.second BSC8.tail) . BSC8.readInt parseChar :: [Char] -> VU.Vector Char parseChar = VU.fromList parse1 :: IO Int parse1 = readLn parse2 :: IO (Int, Int) parse2 = (\vec -> (vec VU.! 0, vec VU.! 1)) . VU.unfoldrN 2 parseInt <$> BSC8.getLine parse3 :: IO (Int, Int, Int) parse3 = (\vec -> (vec VU.! 0, vec VU.! 1, vec VU.! 2)) . VU.unfoldrN 3 parseInt <$> BSC8.getLine parse4 :: IO (Int, Int, Int, Int) parse4 = (\vec -> (vec VU.! 0, vec VU.! 1, vec VU.! 2, vec VU.! 3)) . VU.unfoldrN 4 parseInt <$> BSC8.getLine parseM :: Int -> IO (VU.Vector Int) parseM m = VU.unfoldrN m parseInt <$> BSC8.getLine parseN :: Int -> IO (VU.Vector Int) parseN n = VU.replicateM n parse1 parseNM :: Int -> Int -> IO (V.Vector (VU.Vector Int)) parseNM n m = V.replicateM n $ VU.unfoldrN m parseInt <$> BSC8.getLine fi :: Int -> Integer fi = fromIntegral fI :: Integer -> Int fI = fromInteger gcdInt :: Int -> Int -> Int gcdInt a b = fI $ gcdInteger (fi a) (fi b) gcdextInt :: Int -> Int -> (Int, Int) gcdextInt a b = case gcdExtInteger c d of (# x, y #) -> (fI x, fI y) where c = fromIntegral a d = fromIntegral b lcmInt :: Int -> Int -> Int lcmInt a b = fI $ lcmInteger (fi a) (fi b) sqrInt :: Int -> Int sqrInt = fI . sqrInteger . fi powModInt :: Int -> Int -> Int -> Int powModInt a n mo = fI $ powModInteger (fi a) (fi n) (fi mo) recipModInt :: Int -> Int -> Int recipModInt x n = fI $ recipModInteger (fi x) (fi n) ------------------------------------------------------------------------------- -- FFT ------------------------------------------------------------------------------- dft' :: Double -> VU.Vector (Complex Double) -> VU.Vector (Complex Double) dft' !sign !f | VU.length f == 1 = f | otherwise = let !n = VU.length f !f0 = dft' sign $ VU.map (f VU.!) $ VU.enumFromStepN 0 2 nd2 !f1 = dft' sign $ VU.map (f VU.!) $ VU.enumFromStepN 1 2 nd2 !nd2 = n `div` 2 !zeta = cis (sign * 2 * acos (-1) / fromIntegral n) :: Complex Double in VU.imap (\i z -> (f0 VU.! (i `mod` nd2)) + z * (f1 VU.! (i `mod` nd2))) $ VU.iterateN n (* zeta) 1 {-# INLINE dft' #-} dft, idft :: VU.Vector (Complex Double) -> VU.Vector (Complex Double) dft = dft' 1 idft = dft' (-1) -- FFTによる畳み込み multiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int multiply !f !g = let !len = VU.length f + VU.length g - 1 !sz = fix (\rec !k -> if k >= len then k else rec (k * 2)) 1 !f' = VU.map fromIntegral f VU.++ VU.replicate (sz - VU.length f) (0 :: Complex Double) !g' = VU.map fromIntegral g VU.++ VU.replicate (sz - VU.length g) (0 :: Complex Double) in VU.map (round . magnitude . (/ fromIntegral sz)) . VU.take len . idft $ VU.zipWith (*) (dft f') (dft g') {-# INLINE multiply #-} ------------------------------------------------------------------------------- -- NTT ------------------------------------------------------------------------------- nextPowerOfTwo :: VU.Vector Int -> VU.Vector Int nextPowerOfTwo vs | VU.null vs = VU.singleton 0 | VU.length vs == 1 = vs | n <- unsafeShiftRL (-1) (countLeadingZeros (VU.length vs - 1)) + 1 = vs VU.++ VU.replicate (n - VU.length vs) 0 infixl 8 `shiftRL`, `unsafeShiftRL` shiftRL :: Int -> Int -> Int shiftRL = unsafeShiftRL {-# INLINE shiftRL #-} unsafeShiftRL :: Int -> Int -> Int unsafeShiftRL (I# x#) (I# i#) = I# (uncheckedIShiftRL# x# i#) {-# INLINE unsafeShiftRL #-} bitReverse :: Int -> Int bitReverse = unsafeCoerce @Word64 @Int . step 32 0xffffffff00000000 0x00000000ffffffff . step 16 0xffff0000ffff0000 0x0000ffff0000ffff . step 08 0xff00ff00ff00ff00 0x00ff00ff00ff00ff . step 04 0xf0f0f0f0f0f0f0f0 0x0f0f0f0f0f0f0f0f . step 02 0xcccccccccccccccc 0x3333333333333333 . step 01 0xaaaaaaaaaaaaaaaa 0x5555555555555555 . unsafeCoerce @Int @Word64 where step :: Int -> Word64 -> Word64 -> Word64 -> Word64 step i ml mr = \ !x -> unsafeShiftR (x .&. ml) i .|. unsafeShiftL (x .&. mr) i {-# INLINE step #-} ntt :: Int -> Int -> VU.Vector Int -> VU.Vector Int ntt p g f = runST $ do ff <- VU.unsafeThaw $ VU.unsafeBackpermute f $ VU.generate n ((`unsafeShiftRL` (64 - logN)) . bitReverse) VU.forM_ (VU.iterateN logN (*2) 2) $ \m -> do let !unity = powModInt g (quot (p - 1) m) p let !unities = VU.iterateN (unsafeShiftRL m 1) ((`rem` p) . (* unity)) 1 fix (\loop !k -> when (k < n) $ do flip VU.imapM_ unities $ \j w -> do u <- VUM.unsafeRead ff (k + j) t <- (* w) <$!> VUM.unsafeRead ff (k + j + unsafeShiftRL m 1) VUM.unsafeWrite ff (k + j) $ rem (u + t) p VUM.unsafeWrite ff (k + j + unsafeShiftRL m 1) $ mod (u - t) p loop (k + m) ) 0 VU.unsafeFreeze ff where !n = VU.length f !logN = countTrailingZeros n {-# INLINE ntt #-} intt :: Int -> Int -> VU.Vector Int -> VU.Vector Int intt p g f = VU.map ((`rem` p) . (* n')) $ ntt p (recipModInt g p) f where !n' = recipModInt (VU.length f) p {-# INLINE intt #-} convolute' :: Int -> Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int convolute' p g xs ys = intt p g $ VU.zipWith (\x y -> x * y `rem` p) (ntt p g $ xs VU.++ VU.replicate n 0) (ntt p g $ ys VU.++ VU.replicate n 0) where !n = VU.length xs {-# INLINE convolute' #-} convolute :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int convolute n xs ys = VU.slice 1 (2 * n) $ convolute' 998244353 3 (nextPowerOfTwo $ VU.cons 0 xs) (nextPowerOfTwo $ VU.cons 0 ys) ------------------------------------------------------------------------------- #define MOD 1000000007 modulus :: (Num a) => a modulus = MOD {-# INLINE modulus #-} infixr 8 ^% infixl 7 *%, /% infixl 6 +%, -% (+%) :: Int -> Int -> Int (+%) (I# x#) (I# y#) = case x# +# y# of r# -> I# (r# -# ((r# >=# MOD#) *# MOD#)) {-# INLINE (+%) #-} (-%) :: Int -> Int -> Int (-%) (I# x#) (I# y#) = case x# -# y# of r# -> I# (r# +# ((r# <# 0#) *# MOD#)) {-# INLINE (-%) #-} (*%) :: Int -> Int -> Int (*%) (I# x#) (I# y#) = I# ((x# *# y#) `remInt#` MOD#) {-# INLINE (*%) #-} (/%) :: Int -> Int -> Int (I# x#) /% (I# y#) = go# y# MOD# 1# 0# where go# a# b# u# v# | isTrue# (b# ># 0#) = case a# `quotInt#` b# of q# -> go# b# (a# -# (q# *# b#)) v# (u# -# (q# *# v#)) | otherwise = I# ((x# *# (u# +# MOD#)) `remInt#` MOD#) {-# INLINE (/%) #-} (^%) :: Int -> Int -> Int (^%) x n | n > 0 = go 1 x n | n== 0 = 1 | otherwise = go 1 ((/%) 1 x) (-n) where go !acc !y !m | m .&. 1 == 0 = go acc (y *% y) (unsafeShiftR m 1) | m == 1 = acc *% y | otherwise = go (acc *% y) (y *% y) (unsafeShiftR (m - 1) 1) {-# INLINE (^%) #-} intMod :: (Integral int) => int -> Mint intMod x = fromIntegral $ mod (fromIntegral x) MOD {-# INLINE intMod #-} intModValidate :: Mint -> Bool intModValidate (Mint x) = 0 <= x && x < MOD {-# INLINE intModValidate #-} newtype Mint = Mint { getMint :: Int } deriving newtype (Eq, Ord, Read) instance Show Mint where show (Mint x) = show x instance Num Mint where (+) = coerce (+%) (-) = coerce (-%) (*) = coerce (*%) abs = id signum = const (Mint 1) fromInteger x = coerce @Int @Mint . fromInteger $ mod x MOD instance Real Mint where toRational (Mint x) = toRational x instance Bounded Mint where minBound = Mint 0 maxBound = Mint $ MOD - 1 instance Enum Mint where toEnum = intMod fromEnum = coerce instance Fractional Mint where (Mint x) / (Mint y) = Mint (x /% y) fromRational q = fromInteger (numerator q) / fromInteger (denominator q) instance Integral Mint where quotRem x y = (x / y, x - x / y * y) toInteger = coerce (toInteger @Int) newtype instance VUM.MVector s Mint = MV_Mint (VUM.MVector s Int) newtype instance VU.Vector Mint = V_Mint (VU.Vector Int) instance VU.Unbox Mint instance VGM.MVector VUM.MVector Mint where {-# INLINE basicLength #-} basicLength (MV_Mint v) = VGM.basicLength v {-# INLINE basicUnsafeSlice #-} basicUnsafeSlice i n (MV_Mint v) = MV_Mint $ VGM.basicUnsafeSlice i n v {-# INLINE basicOverlaps #-} basicOverlaps (MV_Mint v1) (MV_Mint v2) = VGM.basicOverlaps v1 v2 {-# INLINE basicUnsafeNew #-} basicUnsafeNew n = MV_Mint `fmap` VGM.basicUnsafeNew n {-# INLINE basicInitialize #-} basicInitialize (MV_Mint v) = VGM.basicInitialize v {-# INLINE basicUnsafeReplicate #-} basicUnsafeReplicate n x = MV_Mint `fmap` VGM.basicUnsafeReplicate n (coerce x) {-# INLINE basicUnsafeRead #-} basicUnsafeRead (MV_Mint v) i = coerce `fmap` VGM.basicUnsafeRead v i {-# INLINE basicUnsafeWrite #-} basicUnsafeWrite (MV_Mint v) i x = VGM.basicUnsafeWrite v i (coerce x) {-# INLINE basicClear #-} basicClear (MV_Mint v) = VGM.basicClear v {-# INLINE basicSet #-} basicSet (MV_Mint v) x = VGM.basicSet v (coerce x) {-# INLINE basicUnsafeCopy #-} basicUnsafeCopy (MV_Mint v1) (MV_Mint v2) = VGM.basicUnsafeCopy v1 v2 {-# INLINE basicUnsafeMove #-} basicUnsafeMove (MV_Mint v1) (MV_Mint v2) = VGM.basicUnsafeMove v1 v2 {-# INLINE basicUnsafeGrow #-} basicUnsafeGrow (MV_Mint v) n = MV_Mint `fmap` VGM.basicUnsafeGrow v n instance VG.Vector VU.Vector Mint where {-# INLINE basicUnsafeFreeze #-} basicUnsafeFreeze (MV_Mint v) = V_Mint `fmap` VG.basicUnsafeFreeze v {-# INLINE basicUnsafeThaw #-} basicUnsafeThaw (V_Mint v) = MV_Mint `fmap` VG.basicUnsafeThaw v {-# INLINE basicLength #-} basicLength (V_Mint v) = VG.basicLength v {-# INLINE basicUnsafeSlice #-} basicUnsafeSlice i n (V_Mint v) = V_Mint $ VG.basicUnsafeSlice i n v {-# INLINE basicUnsafeIndexM #-} basicUnsafeIndexM (V_Mint v) i = coerce `fmap` VG.basicUnsafeIndexM v i basicUnsafeCopy (MV_Mint mv) (V_Mint v) = VG.basicUnsafeCopy mv v {-# INLINE elemseq #-} elemseq _ = seq fact :: Int -> Mint fact = VU.unsafeIndex factCache {-# INLINE fact #-} recipFact :: Int -> Mint recipFact = VU.unsafeIndex recipFactCache {-# INLINE recipFact #-} invFact :: Int -> Mint invFact x = recipFact x * fact (x - 1) {-# INLINE invFact #-} nPk :: Int -> Int -> Mint nPk n k | n < k = 0 | otherwise = fact n * recipFact k {-# INLINE nPk #-} nCk :: Int -> Int -> Mint nCk n k | n < k = Mint 0 | otherwise = fact n * recipFact (n - k) * recipFact k {-# INLINE nCk #-} nHk :: Int -> Int -> Mint nHk n k | n < 0 || k < 0 = Mint 0 | k == 0 = Mint 1 | otherwise = nCk (n + k - 1) k {-# INLINE nHk #-} #define FACT_CACHE_SIZE 200000 factCacheSize :: Int factCacheSize = min (modulus - 1) FACT_CACHE_SIZE {-# INLINE factCacheSize #-} factCache :: VU.Vector Mint factCache = VU.scanl' (\ x y -> x * coerce y) (1 :: Mint) $ VU.generate factCacheSize (+ 1) {-# NOINLINE factCache #-} recipFactCache :: VU.Vector Mint recipFactCache = VU.scanr' ((*) . coerce) (1 / factCache VU.! factCacheSize) $ VU.generate factCacheSize (+ 1) {-# NOINLINE recipFactCache #-} powN :: Int -> VU.Vector Mint powN n = VU.scanl' (\acc x -> acc * intMod x) 1 $ VU.replicate factCacheSize n stirling2nd :: Int -> Int -> Mint stirling2nd n k = (snd ret) * (recipFact k) where xs = zip [0..] $ map f [0..k] f i = (Mint $ powModInt i n modulus) * (nCk k i) ret = List.foldl1' g xs g :: (Int, Mint) -> (Int, Mint) -> (Int, Mint) g acc x | odd (k - fst x) = (fst acc + 1, snd acc - snd x) | otherwise = (fst acc + 1, snd acc + snd x) lagrangePolynomial :: Int -> VU.Vector Mint -> Int -> IO Mint lagrangePolynomial k xs t = do dp <- VUM.replicate (k + 1) (Mint 1) pd <- VUM.replicate (k + 1) (Mint 1) forM_ [0 .. k - 1] $ \i -> do a <- VUM.read dp i VUM.write dp (i + 1) (intMod (t - i) * a) forM_ [k, k - 1 .. 1] $ \i -> do b <- VUM.read pd i VUM.write pd (i - 1) (intMod (t - i) * b) ps <- VU.unsafeFreeze dp qs <- VU.unsafeFreeze pd lp xs ps qs (Mint 0) 0 k where lp :: VU.Vector Mint -> VU.Vector Mint -> VU.Vector Mint -> Mint -> Int -> Int -> IO Mint lp as ps qs ans i j | i > j = return ans | odd (j - i) = lp as ps qs (ans - tmp) (i + 1) j | otherwise = lp as ps qs (ans + tmp) (i + 1) j where tmp = (as VU.! i) * (ps VU.! i) * (qs VU.! i) * (recipFact i) * (recipFact (j - i)) logMod :: Int -> Int -> Int -> Maybe Int logMod a b p = go 0 b where !sp = ceiling . sqrt . fromIntegral $ p !g = powModInt a (-sp) p babystep x = a * x `rem` p giantstep x = g * x `rem` p table :: IntMap Int !table = IntMap.fromList $ zip (iterate babystep 1) [0..sp-1] go !i !x | i < sp = case IntMap.lookup x table of Just j -> Just $! i * sp + j Nothing -> go (i + 1) $ giantstep x | otherwise = Nothing