AtCoder Typical Contest 001

Submission #421042

Source codeソースコード

{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NumDecimals #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExplicitForAll #-}

import Control.Applicative
import Data.Bits
import qualified Data.ByteString.Char8 as BS
import Data.Complex
import qualified Data.Vector.Unboxed as U
import GHC.Exts (Word)

main :: IO ()
main = do
  n <- getInt
  abbs <- U.replicateM n getInt2
  U.mapM_ print $ solve abbs

solve :: U.Vector (Int, Int) -> U.Vector Int
solve (U.unzip -> (as, bs)) =
  U.map (round . realPart) $
  U.take (2 * U.length as) $
  U.drop 1 $
  convolve (toC as) (toC bs)
  where
    toC = U.map fromIntegral . U.cons 0

convolve :: U.Vector (Complex Double) -> U.Vector (Complex Double) -> U.Vector (Complex Double)
convolve as bs = U.map (*rt) $ U.take len $ ifft $ U.zipWith (*) (fft as') (fft bs')
  where
    as' = extend n as
    bs' = extend n bs
    len = U.length as + U.length bs - 1
    n = 1 `shiftL` ((1+) $ bitScanReverse $ fromIntegral len)
    rt = recip $ fromIntegral n

    extend k xs = xs U.++ U.replicate (k - U.length xs) 0

fft :: U.Vector (Complex Double) -> U.Vector (Complex Double)
fft vec
  | U.length vec == 1 = vec
fft vec = U.generate n $ \i ->
  vec1 U.! (mod i nh) +
  vec2 U.! (mod i nh) * cis (fromIntegral i / fromIntegral n * 2 * pi)
  where
    vec1 = fft $ U.generate nh $ \i -> vec U.! (2 * i)
    vec2 = fft $ U.generate nh $ \i -> vec U.! (2 * i + 1)
    !n = U.length vec
    !nh = div n 2

ifft :: U.Vector (Complex Double) -> U.Vector (Complex Double)
ifft vec
  | U.length vec == 1 = vec
ifft vec = U.generate n $ \i ->
  vec1 U.! (mod i nh) +
  vec2 U.! (mod i nh) * cis (-fromIntegral i / fromIntegral n * 2 * pi)
  where
    vec1 = ifft $ U.generate nh $ \i -> vec U.! (2 * i)
    vec2 = ifft $ U.generate nh $ \i -> vec U.! (2 * i + 1)
    !n = U.length vec
    !nh = div n 2

----------------------------------------------------------------------------
-- IO

getInt :: IO Int
getInt = readInt <$> BS.getLine

readInt :: BS.ByteString -> Int
readInt s = case BS.readInt s of
  Just (r, "") -> r
  _ -> error $ "not an integer: " ++ show s

getInt2 :: IO (Int, Int)
getInt2 = readInt2 <$> BS.getLine

readInt2 :: BS.ByteString -> (Int, Int)
readInt2 s = (v0, v1)
  where
    !v = readIntsN 2 s
    !v0 = U.unsafeIndex v 0
    !v1 = U.unsafeIndex v 1

readIntsN :: Int -> BS.ByteString -> U.Vector Int
readIntsN n s0
  | U.length vec == n = vec
  | otherwise = error $ "readIntsN: expecting " ++ show n
    ++ " ints but got " ++ show (U.length vec)
  where
    vec = U.unfoldrN (n+1) step s0
    step (BS.dropWhile (==' ') -> s)
      | s == "" = Nothing
      | Just (v, r) <- BS.readInt s = Just (v, r)
      | otherwise = error $ "not an integer: " ++ show s

----------------------------------------------------------------------------
-- Util

-- | Returns the position of the highest 1 in @w@. If @w@ is 0, returns @0@.
bitScanReverse :: Word -> Int
bitScanReverse w0 = snd $ f 1 $ f 2 $ f 4 $ f 8 $ f 16 $ f 32 (w0, 0)
  where
    f k (!w, !acc)
      | w .&. mask /= 0 = (w `shiftR` k, acc + k)
      | otherwise = (w, acc)
      where
        mask = complement 0 `shiftL` k
    {-# INLINE f #-}

Submission

Task問題 C - 高速フーリエ変換
User nameユーザ名 mkotha
Created time投稿日時
Language言語 Haskell (Haskell Platform 2014.2.0.0)
Status状態 AC
Score得点 100
Source lengthソースコード長 3663 Byte
File nameファイル名
Exec time実行時間 4556 ms
Memory usageメモリ使用量 63768 KB

Test case

Set

Set name Score得点 / Max score Cases
Sample - 00_sample_01
All 100 / 100 00_sample_01,01_00_01,01_01_19,01_02_31,01_03_22,01_04_31,01_05_40,01_06_15,01_07_39,01_08_28,01_09_30,01_10_23,01_11_33,01_12_11,01_13_28,01_14_41,01_15_26,01_16_49,01_17_34,01_18_02,01_19_33,01_20_29,02_00_51254,02_01_82431,02_02_17056,02_03_34866,02_04_6779,02_05_65534,02_06_65535,02_07_65536,02_08_65537,02_09_65538,02_10_100000

Test case

Case name Status状態 Exec time実行時間 Memory usageメモリ使用量
00_sample_01 AC 29 ms 1432 KB
01_00_01 AC 29 ms 1420 KB
01_01_19 AC 31 ms 1692 KB
01_02_31 AC 29 ms 1688 KB
01_03_22 AC 29 ms 1684 KB
01_04_31 AC 29 ms 1696 KB
01_05_40 AC 32 ms 1816 KB
01_06_15 AC 30 ms 1480 KB
01_07_39 AC 30 ms 1944 KB
01_08_28 AC 29 ms 1688 KB
01_09_30 AC 30 ms 1684 KB
01_10_23 AC 28 ms 1720 KB
01_11_33 AC 31 ms 1892 KB
01_12_11 AC 29 ms 1564 KB
01_13_28 AC 32 ms 1808 KB
01_14_41 AC 31 ms 1896 KB
01_15_26 AC 32 ms 1724 KB
01_16_49 AC 32 ms 1904 KB
01_17_34 AC 31 ms 1908 KB
01_18_02 AC 29 ms 1432 KB
01_19_33 AC 30 ms 1944 KB
01_20_29 AC 30 ms 1692 KB
02_00_51254 AC 2212 ms 35220 KB
02_01_82431 AC 4501 ms 63760 KB
02_02_17056 AC 1057 ms 17684 KB
02_03_34866 AC 2163 ms 34004 KB
02_04_6779 AC 273 ms 5396 KB
02_05_65534 AC 2249 ms 37400 KB
02_06_65535 AC 2229 ms 36056 KB
02_07_65536 AC 4508 ms 62676 KB
02_08_65537 AC 4556 ms 62712 KB
02_09_65538 AC 4503 ms 62752 KB
02_10_100000 AC 4550 ms 63768 KB