Submission #62830005


Source Code Expand

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

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

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

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

abc392g :: Int -> [Int] -> Int
abc392g _n ss = sum $ zipWith f ps (olol cs)
  where
    n = maximum ss
    arr = accumArray (+) 0 (0, n) [(s,1) | s <- ss] :: UArray Int Int
    ps = elems arr
    cs = convolution (succ n) ps (succ n) ps
    f 1 c = div (pred c) 2
    f 0 _ = 0
-- 偶数番目を取り出す
    olol (x:xs) = x : olol (drop 1 xs)
    olol _ = []

convolution :: Int -> [Int] -> Int -> [Int] -> [Int]
convolution na as nb bs = runST $
  do
    av <- MUV.new n :: ST s (MUV.MVector s Int)
    zipWithM_ (MUV.write av) [0 ..] as
    bv <- MUV.new n :: ST s (MUV.MVector s Int)
    zipWithM_ (MUV.write bv) [0 ..] bs
    fft av root
    fft bv root
    forM_ [0 .. n1] (\i -> do
      b <- MUV.read bv i
      MUV.modify av (mul b) i
      )
    fft av invroot
    forM [0 .. na + pred nb] (fmap (mul divn) . MUV.read av)
  where
-- 定数
    p = 998244353
    root = 3
    invroot = modRecip root
-- pを底とするモジュロ演算
    reg x = mod x p
    add x y = reg $ x + y
    mul x y = reg $ x * y
    modRecip a = reg $ fst $ loop a p
      where
        loop a 0 = (1, 0)
        loop a b = let (y, x) = loop b (mod a b)
                   in  (x, y - div a b * x)
    pow b e = loop e (reg b) 1
      where
        loop 0 _ acc = acc
        loop k c acc = loop (div k 2) (mul c c) $ if odd k then mul acc c else acc
-- 本体
    (n, w) = (n, length xs)
      where
        (xs, n:_) = span ((na + nb) >) $ iterate (2 *) 1
    n1 = pred n
    divn = modRecip n
-- 下請け
-- ビット反転
    bitReverse :: Int -> Int
    bitReverse j = sum [bit i | i <- [0 .. pred w], testBit j (pred w - i)]
-- 高速フーリエ変換
    fft a base =
      do
-- 配列ビット反転
        forM_ [0 .. n1] (\i -> do
          let j = bitReverse i
          when (i < j) (MUV.swap a i j)
          )
        forM_ (zip omega2N $ iterate (2 *) 1) (\(omega, dist) ->
          forM_ [0, dist + dist .. pred n] (\b -> do
            let (omegas1, omegas2) = splitAt dist $ iterate (mul omega) 1
            forM_ (zip3 omegas1 omegas2 [b ..]) (\(o1, o2, i) -> do
              f_even <- MUV.read a i
              f_odd  <- MUV.read a (i + dist)
              MUV.write a  i         (reg $ f_even + o1 * f_odd)
              MUV.write a (i + dist) (reg $ f_even + o2 * f_odd)
              )
            )
          )
      where
        omegaN = pow base (div (pred p) n)
        omega2N = reverse $ take w $ iterate (\x -> mul x x) omegaN

Submission Info

Submission Time
Task G - Fine Triplets
User joetheshootingst
Language Haskell (GHC 9.4.5)
Score 0
Code Size 2981 Byte
Status TLE
Exec Time 3340 ms
Memory 515776 KiB

Compile Error

app/Main.hs:28:5: warning: [-Wincomplete-patterns]
    Pattern match(es) are non-exhaustive
    In an equation for ‘f’:
        Patterns of type ‘a’, ‘a’ not matched:
            p _ where p is not one of {1, 0}
   |
28 |     f 1 c = div (pred c) 2
   |     ^^^^^^^^^^^^^^^^^^^^^^...

app/Main.hs:56:5: warning: [-Wunused-local-binds]
    Defined but not used: ‘add’
   |
56 |     add x y = reg $ x + y
   |     ^^^

app/Main.hs:60:14: warning: [-Wunused-matches]
    Defined but not used: ‘a’
   |
60 |         loop a 0 = (1, 0)
   |              ^

app/Main.hs:60:14: warning: [-Wname-shadowing]
    This binding for ‘a’ shadows the existing binding
      bound at app/Main.hs:58:14
   |
60 |         loop a 0 = (1, 0)
   |              ^

app/Main.hs:61:14: warning: [-Wname-shadowing]
    This binding for ‘a’ shadows the existing binding
      bound at app/Main.hs:58:14
   |
61 |         loop a b = let (y, x) = loop b (mod a b)
   |              ^

app/Main.hs:70:9: warning: [-Wincomplete-uni-patterns]
    Pattern match(es) are non-exhaustive
    In a pattern binding:
        Patterns of type ‘([Int], [Int])’ not matched: (_, [])
   |
70 |         (xs, n:_) = span ((na + nb) >) $ iterate (2 *) 1
   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

app/Main.hs:70:14: warning: [-Wname-shadowing]
    This binding for ‘n’ shadows the existing binding
      bound at app/Main.hs:68:6
   |
70 |         (xs, n:_) = span ((na + nb) >) $ iterate (2 *) 1
   |              ^

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 0 / 600
Status
AC × 3
AC × 14
TLE × 29
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt
All sample_01.txt, sample_02.txt, sample_03.txt, test_01.txt, test_02.txt, test_03.txt, test_04.txt, test_05.txt, test_06.txt, test_07.txt, test_08.txt, test_09.txt, test_10.txt, test_11.txt, test_12.txt, test_13.txt, test_14.txt, test_15.txt, test_16.txt, test_17.txt, test_18.txt, test_19.txt, test_20.txt, test_21.txt, test_22.txt, test_23.txt, test_24.txt, test_25.txt, test_26.txt, test_27.txt, test_28.txt, test_29.txt, test_30.txt, test_31.txt, test_32.txt, test_33.txt, test_34.txt, test_35.txt, test_36.txt, test_37.txt, test_38.txt, test_39.txt, test_40.txt
Case Name Status Exec Time Memory
sample_01.txt AC 2 ms 6956 KiB
sample_02.txt AC 1619 ms 255360 KiB
sample_03.txt AC 2 ms 7020 KiB
test_01.txt AC 2 ms 6804 KiB
test_02.txt TLE 3340 ms 515540 KiB
test_03.txt AC 2 ms 6872 KiB
test_04.txt AC 2 ms 6920 KiB
test_05.txt TLE 3340 ms 515432 KiB
test_06.txt TLE 3340 ms 515320 KiB
test_07.txt TLE 3340 ms 515292 KiB
test_08.txt TLE 3340 ms 515396 KiB
test_09.txt AC 2 ms 7332 KiB
test_10.txt TLE 3339 ms 512260 KiB
test_11.txt AC 2 ms 7356 KiB
test_12.txt TLE 3339 ms 500140 KiB
test_13.txt AC 5 ms 11056 KiB
test_14.txt TLE 3340 ms 514348 KiB
test_15.txt AC 5 ms 11100 KiB
test_16.txt TLE 3339 ms 501060 KiB
test_17.txt AC 31 ms 15928 KiB
test_18.txt TLE 3340 ms 515548 KiB
test_19.txt AC 31 ms 16052 KiB
test_20.txt TLE 3339 ms 511376 KiB
test_21.txt AC 344 ms 73444 KiB
test_22.txt TLE 3340 ms 515516 KiB
test_23.txt AC 384 ms 71064 KiB
test_24.txt TLE 3340 ms 515516 KiB
test_25.txt TLE 3340 ms 514700 KiB
test_26.txt TLE 3340 ms 515664 KiB
test_27.txt TLE 3339 ms 515548 KiB
test_28.txt TLE 3340 ms 515616 KiB
test_29.txt TLE 3339 ms 515656 KiB
test_30.txt TLE 3339 ms 514800 KiB
test_31.txt TLE 3339 ms 515564 KiB
test_32.txt TLE 3339 ms 515776 KiB
test_33.txt TLE 3339 ms 515556 KiB
test_34.txt TLE 3339 ms 515636 KiB
test_35.txt TLE 3338 ms 515636 KiB
test_36.txt TLE 3339 ms 515772 KiB
test_37.txt TLE 3339 ms 515720 KiB
test_38.txt TLE 3339 ms 515620 KiB
test_39.txt TLE 3338 ms 515564 KiB
test_40.txt TLE 3338 ms 515552 KiB