Codeforces 113C(114E)
前回参加したCodeforces #86 Div.2のE(＝Div.1のC)
Double HappinessをHaskellで。
問題は区間 [l,r] に含まれる
・素数
かつ
・２つの平方数の和
であるような自然数の数を出力せよとのこと。
問題は単純なんで正解を出力するだけならいくらでも書きようあるのですが、時間制限の5secがネック。
1~300000000までの区間を1000000ぐらいに区切って予めその区間に置ける答えを求めておき、
[200,3500000] = [200,1000000]
+ [1000000,2000000](計算済)
+ [2000000,3000000](計算済)
+ [3000000,3500000]
などと分割してやるのが正解だったようですが、（気づかなかったので）安直に篩をつかって間に合わせてみました。
ものすごく厳しかった。
まず、
のページの一番下に書いてあるとおり、奇素数nが２つの平方数の和で表される必要十分条件は「mod n 4==1」です。
区間に2が含まれる時だけ注意すればこれは問題ないですね。
なので方針は、細かいところを除けば
(1)[1,r]を篩う。
(2)[l,r]から素数かつ4で割ったあまりが1である数をカウント
となります。
ただこれだと間に合わないです。微妙に。
そこで、篩を作る時のことを考えると、
そもそも(2)でカウントするアルゴリズムを
(2')begin := (l以上で最小の「4n+1」型整数)
として i = begin,begin+4,...,l についてiが素数かどうかを調べる。
とすれば、そもそも偶数については素数かどうかの判定が起こらないので、偶数を篩う必要がなくなります。
しかも(2')の方が(2)よりそもそも速いですしね。
さて、アルゴリズム的にはこれでいいのですが、あとはHaskellの動作をいかに早くするかの勝負です。
例えば、UArrayへのアクセスはunsafeAt関数のほうが微妙に速いです。
篩を作る関数sieveも「Haskellでエラトステネスの篩」：http://d.hatena.ne.jp/g940425/20110827/1314442246 に追記した関数を用います。かなり高速化できたと思いましたが、これを使ってギリギリでした。
コードは以下。
{-# OPTIONS_GHC -O2 #-} import Data.Array.ST import Data.Array.Unboxed import Data.Array.Base (unsafeWrite,unsafeRead,unsafeAt) import Monad class Scan a where scan' :: String -> a instance Scan Int where scan' = read instance (Scan a,Scan b) => Scan (a,b) where scan' x = (\(x:y:_) -> (scan' x,scan' y)) (words x) scan :: (Scan a) => IO a scan = getLine>>=(return.scan') sieve :: Int -> UArray Int Bool sieve n = runSTUArray $ do t <- newArray (0,n) True unsafeWrite t 1 False let sqn = (floor.sqrt.fromIntegral) (n::Int) mapM_ (\i -> unsafeRead t i >>= (flip when) (mapM_ (\j -> (unsafeWrite t j False)) [i*i,i*(i+2)..n])) [3,5..sqn] return t solve l r = rec beg 0 where sve = sieve r beg = case mod l 4 of 0 -> l+1 1 -> l 2 -> l+3 3 -> l+2 rec i n |i>r = n + (if l<=2&&r>=2 then 1 else 0) |unsafeAt sve i = rec (i+4) (n+1) |otherwise = rec (i+4) n main = do (l,r) <- scan :: IO (Int,Int) print (solve l r)
これならなんとか間に合います。
でもコレ時間内に書けと言われたら厳しいものがありますね……精進せねば。
ではまた。
