AtCoder Typical Contest 001

Submission #420844

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.cons 0 $
  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) $ ifft $ U.zipWith (*) (fft as') (fft bs')
  where
    as' = extend n as
    bs' = extend n bs
    n = 1 `shiftL` (bitScanReverse $ fromIntegral $ U.length as + U.length bs)
    rt = recip $ fromIntegral n

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

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状態 WA
Score得点 0
Source lengthソースコード長 3627 Byte
File nameファイル名
Exec time実行時間 ms
Memory usageメモリ使用量 -

Test case

Set

Set name Score得点 / Max score Cases
Sample - 00_sample_01
All 0 / 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 95 ms 1408 KB
01_00_01 WA
01_01_19 WA
01_02_31 WA
01_03_22 WA
01_04_31 WA
01_05_40 WA
01_06_15 WA
01_07_39 WA
01_08_28 WA
01_09_30 WA
01_10_23 WA
01_11_33 WA
01_12_11 WA
01_13_28 WA
01_14_41 WA
01_15_26 WA
01_16_49 WA
01_17_34 WA
01_18_02 AC 28 ms 1372 KB
01_19_33 WA
01_20_29 WA
02_00_51254 WA
02_01_82431 WA
02_02_17056 WA
02_03_34866 WA
02_04_6779 WA
02_05_65534 WA
02_06_65535 WA
02_07_65536 AC 2248 ms 38172 KB
02_08_65537 WA
02_09_65538 WA
02_10_100000 WA