AtCoder Typical Contest 001

Submission #421898

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 Control.Monad
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.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)

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

ifft :: U.Vector (Complex Double) -> U.Vector (Complex Double)
ifft vec = U.map ((*r) . conjugate) $ fft $ U.map conjugate vec
  where
    r = recip $ fromIntegral $ U.length vec

fft :: U.Vector (Complex Double) -> U.Vector (Complex Double)
fft vec
  | U.length vec == 1 = vec
fft !vec = U.zipWith (+) vec1 vec2 U.++ U.zipWith (+) vec1 (U.map negate vec2)
  where
    !rot = cis $ 2 * pi / fromIntegral n
    !vec1 = fft $ U.generate nh $ \i -> vec `U.unsafeIndex` (2 * i)
    !vec2 = U.zipWith (*) (U.iterateN nh (*rot) 1) $
      fft $ U.generate nh $ \i -> vec `U.unsafeIndex` (2 * i + 1)
    !n = U.length vec
    !nh = n `shiftR` 1

----------------------------------------------------------------------------
-- 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ソースコード長 3497 Byte
File nameファイル名
Exec time実行時間 1359 ms
Memory usageメモリ使用量 47184 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 30 ms 1228 KB
01_00_01 AC 29 ms 1232 KB
01_01_19 AC 30 ms 1540 KB
01_02_31 AC 30 ms 1652 KB
01_03_22 AC 31 ms 1712 KB
01_04_31 AC 29 ms 1688 KB
01_05_40 AC 30 ms 1748 KB
01_06_15 AC 30 ms 1364 KB
01_07_39 AC 30 ms 1812 KB
01_08_28 AC 30 ms 1620 KB
01_09_30 AC 29 ms 1620 KB
01_10_23 AC 29 ms 1620 KB
01_11_33 AC 30 ms 1748 KB
01_12_11 AC 29 ms 1408 KB
01_13_28 AC 29 ms 1624 KB
01_14_41 AC 32 ms 1744 KB
01_15_26 AC 30 ms 1624 KB
01_16_49 AC 33 ms 1744 KB
01_17_34 AC 32 ms 1752 KB
01_18_02 AC 30 ms 1236 KB
01_19_33 AC 33 ms 1812 KB
01_20_29 AC 31 ms 1632 KB
02_00_51254 AC 704 ms 22612 KB
02_01_82431 AC 1331 ms 46292 KB
02_02_17056 AC 336 ms 12984 KB
02_03_34866 AC 657 ms 24636 KB
02_04_6779 AC 110 ms 4944 KB
02_05_65534 AC 729 ms 25688 KB
02_06_65535 AC 734 ms 25684 KB
02_07_65536 AC 1292 ms 44944 KB
02_08_65537 AC 1290 ms 47184 KB
02_09_65538 AC 1282 ms 44248 KB
02_10_100000 AC 1359 ms 44184 KB