結果

問題 No.206 数の積集合を求めるクエリ
ユーザー かりあげクンかりあげクン
提出日時 2020-09-02 21:46:39
言語 Haskell
(9.8.2)
結果
CE  
(最新)
AC  
(最初)
実行時間 -
コード長 14,960 bytes
コンパイル時間 272 ms
コンパイル使用メモリ 168,704 KB
最終ジャッジ日時 2024-04-27 03:29:47
合計ジャッジ時間 914 ms
ジャッジサーバーID
(参考情報)
judge1 / judge5
このコードへのチャレンジ
(要ログイン)
コンパイルエラー時のメッセージ・ソースコードは、提出者また管理者しか表示できないようにしております。(リジャッジ後のコンパイルエラーは公開されます)
ただし、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

ソースコード

diff #

{-# 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

  let cx = convolute 262144 as bs
  putStr $ unlines $ map show $ VU.toList $ VU.drop n $ VU.take (n + q) cx

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)

-------------------------------------------------------------------------------
-- 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
0