Haskellで二分ヒープ(優先度付きキュー)実装してみた
Haskellには、C++で言う所のpriority_queueにあたるモジュールがないので、Codeforcesでこれが必要な問題は解けないという問題を抱えていた（幸い、これまで自分が出てきたコンテストでそういう問題には当たらなかった、はずです）ので、自前で実装してみました。
実装したのは
型
| Heap a | a型の変数を持つヒープ |
関数
| 関数名 | 型 | 機能 |
|---|---|---|
| make_heap | (a -> a -> Bool) -> [a] -> Int -> Heap a | 順序付けする関数（後述）と要素のリスト、ヒープの最大サイズを受け取ってヒープを構築 |
| pop_heap | Heap a -> Maybe (a, Heap a) | 最大（最小）値を取り出す。ヒープが空の時はNothingを返す |
| push_heap | Heap a -> a -> Heap a | ヒープに要素を挿入 |
です。他にも細々とした補助関数は定義しましたが。
コードは以下。
なお、ヒープが本来要求される計算量
make : O(n)
pop : O(log n)
push : O(log n)
を達成するにはGHCで-Oあるいは-O2オプションをつけてコンパイルすることが必要です。
import Data.Array.ST import Control.Monad.ST import Data.Array import Data.List import Monad data Heap a = H {size::Int,ary::Array Int a,pol:: a -> a -> Bool} instance (Show a) => Show (Heap a) where show H {size=n,ary=a} = "Heap [" ++ (intercalate "," (map (show.(a!)) [1..n]))++"]" build_heap :: (a -> a -> Bool) -> [a] -> Int -> Array Int a build_heap policy l n = runSTArray $ do let len = length l ary <- newListArray (1,n) l let swap i j ei ej = writeArray ary i ej >> writeArray ary j ei let heapity_rec i = let left = 2*i;right = 2*i+1 in do elem <- readArray ary i left_elem <- if left<=len then readArray ary left else return elem right_elem <- if right<=len then readArray ary right else return elem (if left<=len&&policy left_elem elem then if right<=len&&policy right_elem left_elem then swap i right elem right_elem>>heapity_rec right else swap i left elem left_elem>>heapity_rec left else when (right<=len&&policy right_elem elem) (swap i right elem right_elem>>heapity_rec right)) let beg = div len 2 in mapM_ heapity_rec [beg,beg-1..1] return ary make_heap :: (a -> a -> Bool) -> [a] -> Int -> Heap a make_heap policy l n = H {size=length l,ary=(build_heap policy l n),pol=policy} top_heap :: Heap a -> a top_heap = (!1).ary pop_heap :: Heap a -> Maybe (a, Heap a) pop_heap heap@H{size=len,ary=a,pol=policy} = if len==0 then Nothing else Just (top_heap heap, heap{size=len-1,ary=heapity}) where heapity = runSTArray $ do ary <- unsafeThaw a readArray ary len >>= writeArray ary 1 let swap i j ei ej = writeArray ary i ej >> writeArray ary j ei let heapity_rec i = let left = 2*i;right = 2*i+1 in do elem <- readArray ary i left_elem <- if left<=len then readArray ary left else return elem right_elem <- if right<=len then readArray ary right else return elem (if left<=len&&policy left_elem elem then if right<=len&&policy right_elem left_elem then swap i right elem right_elem>>heapity_rec right else swap i left elem left_elem>>heapity_rec left else when (right<=len&&policy right_elem elem) (swap i right elem right_elem>>heapity_rec right)) heapity_rec 1 return ary push_heap :: Heap a -> a -> Heap a push_heap heap@H{size=len,ary=a,pol=policy} x = H{size=len+1,ary=heapity,pol=policy} where heapity = runSTArray $ do ary <- unsafeThaw a writeArray ary (len+1) x let swap i j ei ej = writeArray ary i ej >> writeArray ary j ei let heapity_rec i = do let parent = div i 2 elem <- readArray ary i parent_elem <- readArray ary parent when (i>1&&policy elem parent_elem) (swap i parent elem parent_elem >> when (parent>1) (heapity_rec parent)) when (len>0) (heapity_rec (len+1)) return ary
実装はよく見る配列による実装です。
二分木型を作ってその上で実装することも考えたのですが、
資料の多さからとりあえずこの方法を取りました。
Heap型には本体となる配列とヒープ末尾の位置を持たせています。
max・minヒープどちらにでも対応するために、
make_heapには要素を順序付けるための関数（述語）をもたせる仕様とし、
make_heap (<) ……でmin-ヒープ
make_heap (>) ……でmax-ヒープ
をそれぞれ構成できます。
STモナドでの配列操作が増えた結果かなり煩雑なコードになってしまったのが惜しいところですが、速度はまあまあです。
これを優先度付きキューとして使った蟻本の問題でもそのうち解いてみようと思います。
それではﾉｼ
TopCoder:SRM 518
SRMに参戦。
いい結果でしたが、妙にEasy・Mediumが簡単だった？ような気がします。早解き出来るかが勝負の分かれ目になりそうです。
Easy: 208.12
Medium: 437.90
Hard: Opened
646.02pts
77th
Rating :1469 -> 1644
なんにせよ黄色には復帰＆自己ベスト更新しました。
よかったです。
自分の解法をちょろっと。
コードは今回は載せません。
Easy
(1)str = ""
(2)sの中で最も大きい文字cを選んでstrの末尾に加える。
(3)s = （sのcより右側の部分文字列）
(4)s=""でなければ(2)に戻る。
これでOKです。
Medium
階差数列で解いてる人が多かったようですが、こんな単純なアルゴリズムで通ります。局所改善法？
(1)ans=0
(2)aがconvexか見る。convexならansを返す。
(3)for文でaを先頭から見ていき、
a[i]*2 <= a[i-1] + a[i-1]
を満たさなければa[i]の値を更新。
(4)(2)に戻る。
これでOKです。
なぜなら、各操作の時点では必要最小限の値しかa[i]を減少させていないので、この操作でa[i]の値を真の解より下げすぎることがないからです。
しかも、この操作を続けていけば必ずconvexになるのは自明なので、このループは必ず止まります。
問題は時間ですが、
1,1000000000,1000000000, .. , 1000000000
でも普通に余裕をもって解を出力してくれました。
今回は以上です。
Hardは無理でした。Nimで後攻が必勝になるための条件は知っていましたが、それでもKの値が大きすぎて全パターンを調べ上げる効率的な方法を思いつくことができませんでした。
それではまた。
このまませっかく上がったレート＆黄色を維持できるようにがんばります。
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)
これならなんとか間に合います。
でもコレ時間内に書けと言われたら厳しいものがありますね……精進せねば。
ではまた。
TopCoder:SRM 517
参戦してきましたが一問もあってなかったので結果だけorz
250: Failed System Test
600: Opened
1000: UnOpened
Challenge: +50/25
Score: 25pts
Ranking: 581st
Rating: 1512 -> 1469
青に戻ってしまいました……
まだまだ実力が足らぬようです。
水曜にもっかいあるのでリベンジしたいと思います！
では！
Codeforces Beta Round #86 (Div. 2 Only)(参戦報告+A,B,C解法)
Codeforces Beta Round #86 Div.2に参戦。
言語はHaskell。
結果
A:488
B:824
C:1024
D:TLE
E:-2
位
2336 pts
Rating:1573 -> 1659
Eはpretestの時点でTLEが最後まで取れず、DもTLEで落ちました。
でもDiv2内で順位が2ケタ台までいけたし、なにより紫昇格＆Div 1昇格できたので個人的には上々の出来。
通した問題の解法をちょろっと。
A:
kがnの累乗数になってるか確認するだけですね。
{-# LANGUAGE TypeSynonymInstances #-} {-# OPTIONS_GHC -O2 #-} import Debug.Trace import Data.List import Data.Array import Data.Char class Scan a where scan' :: String -> a instance Scan Int where scan' = read instance Scan Char where scan' (x:_) = x instance Scan Float where scan' = read instance Scan Double where scan' = read instance Scan Integer where scan' = read instance Scan String where scan' x = x instance (Scan a,Scan b) => Scan (a,b) where scan' x = (\(x:y:_) -> (scan' x,scan' y)) (words x) instance (Scan a,Scan b,Scan c) => Scan (a,b,c) where scan' x = (\(x:y:z:_) -> (scan' x,scan' y,scan' z)) (words x) instance (Scan a,Scan b,Scan c,Scan d) => Scan (a,b,c,d) where scan' x = (\(w:x:y:z:_) -> (scan' w,scan' x,scan' y,scan' z)) (words x) instance (Scan a,Scan b,Scan c,Scan d,Scan e) => Scan (a,b,c,d,e) where scan' x = (\(v:w:x:y:z:_) -> (scan' v,scan' w,scan' x,scan' y,scan' z)) (words x) class Ans a where showans :: a -> String instance Ans Int where showans x = show x instance Ans Char where showans x = [x] instance Ans Float where showans x = show x instance Ans Double where showans x = show x instance Ans Integer where showans x = show x instance Ans String where showans x = x instance (Ans a, Ans b) => Ans (a,b) where showans (x,y) = showans x ++ " " ++ showans y instance (Ans a, Ans b,Ans c) => Ans (a,b,c) where showans (x,y,z) = showans x ++ " " ++ showans y ++ " " ++ showans z scan :: (Scan a) => IO a scan = getLine>>=(return.scan') scans :: (Scan a) => Int -> IO [a] scans n = if n==0 then return [] else scan>>=(\x->scans (n-1)>>=return.(x:)) scanlist :: (Scan a) => IO [a] scanlist = getLine>>=return.(map scan').words scanlists :: (Scan a) => Int -> IO [[a]] scanlists n = if n==0 then return [] else scanlist>>=(\x->scanlists (n-1)>>=return.(x:)) putAnsLn :: (Ans a) => a -> IO () putAnsLn = putStrLn.showans putAnsLns :: (Ans a) => [a] -> IO () putAnsLns = mapM_ putAnsLn --ここまで入出力ライブラリ solve k n m = if mod n k /= 0 then ["NO"] else if div n k == 1 then "YES":[show m] else solve k (div n k) (m+1) main = do k <- scan :: IO Int n <- scan ::IO Int putAnsLns (solve k n 0)
以下、入出力ライブラリ部分は省きます。
import宣言は書きます。
B:
n≦16と、問題のサイズが小さかったのでDFSで全探索。
リストを先頭から見ていき、「先頭の人をチームに入れた場合の最大値」と「先頭の人をチームに入れなかった場合の最大値」を比べて大きい方を返す関数recでDFSを行ってます。
同時にチームに入れられない２人の組み合わせはMapにリストで格納し、チームに入れた人間と一緒に入れられない人間をリストから外しながら再帰しています。リストの差分を返す演算子(\\)は重複があった場合一つしかリストから除きませんが、問題文に「n人のメンバーに重複はない」とあるので、そもそも重複を考えなくて良いことがわかります。
import Debug.Trace import Data.List import Data.Array import Data.Char import Data.Map rec _ [] l = (length l,l) rec antis (x:xs) l = let (a1,team1) = rec antis xs l in let diff' = (Data.Map.lookup x antis) in let diff = case diff' of Nothing -> [] Just l -> l in let (a2,team2) = rec antis (xs Data.List.\\ diff) (x:l) in if a1>a2 then (a1,team1) else (a2,team2) solve members n anti = let antis = foldl (\mp -> \(na,an) -> insertWith (++) an [na] (insertWith (++) na [an] mp)) empty anti in let (ans,team) = rec antis members [] in if ans==0 then ["0"] else (show ans):(sort team) main = do (n,m) <- scan :: IO (Int,Int) members <- scans n :: IO [String] anti <- scans m :: IO [(String,String)] putAnsLns (solve members n anti)
C:
１単語 or 1文 であるときにYESを出力する問題。１単語の場合を忘れて一回pretestで弾かれました。
HaskellにはData.Listモジュールに、ある文字列の接尾語が指定したsuffixと一致しているかどうか調べるisSuffixOfという関数がありますので、これを使って接尾辞を調べられます。
性はBool値fで区別します。
isAdj f = isSuffixOf (if f then "lios" else "liala") isNoun f = isSuffixOf (if f then "etr" else "etra") isVerb f = isSuffixOf (if f then "initis" else "inites")
文章として適した順番かどうか調べるときは、f=True,Falseについて
(1)dropWhile (isAdj f) (入力)で先頭のAdj部分を切り落とし
(2)isNounで次の一語がNounであるか調べ
(3)all (isVerb f) (残り)で残りが全てVerbであるか調べる。
という手順であっさりと書けます。
import Debug.Trace import Data.List import Data.Array import Data.Char isAdj f = isSuffixOf (if f then "lios" else "liala") isNoun f = isSuffixOf (if f then "etr" else "etra") isVerb f = isSuffixOf (if f then "initis" else "inites") parse f s = let s' = dropWhile (isAdj f) s in case s' of [] -> "NO" (x:xs) -> if isNoun f x && all (isVerb f) xs then "YES" else "NO" validword [] = False validword (x:_) = (isVerb True x) ||(isVerb False x) ||(isNoun True x) ||(isNoun False x) ||(isAdj True x) ||(isAdj False x) solve s = if length s == 1 && validword s then "YES" else if parse True s == "NO" then parse False s else "YES" main = do statement <- scanlist :: IO [String] putAnsLn (solve statement)
そういえば、少しIOライブラリを整理・追加しました。
可変個の引数取れないなどの問題があったので。
あと、mapM_などを用いたりdo記法を省いたりして全体的にコード量を減らしました。
最近ちょっとは読みやすいコードを書けるようになってきた気がします。おそらくまだまだなのでしょうが。
map・fold・filterなど、リストをまとめて処理する関数はなるべく使うようにしています。Haskellerらしくボイラープレートはなるべく排除していきたいです。
ではでは（・ω・）ﾉｼ
Codeforces Beta Round #85 (Div. 2 Only)
Codeforces Beta Round #85(Div.2)に参戦。
言語はHaskell。
結果
A:484(00:08)
B:928(00:18)
C:1302(00:33)
D:(-4)
E:No Submit
Hack : No Hacks
2714pts 150位
Rating:1476 -> 1573
Dは一応提出しました。正解は出力するようですが、安直に書きすぎてTLEしてしまいました。
それでは通った問題の解法を。
以下に載せるコードは入出力周りを省いてあります。
入出力関連の関数ライブラリはこの記事
http://d.hatena.ne.jp/g940425/20110823/1314082061
にあります。
A:
小文字に直す→比べる
それだけ。
小文字に直す関数としてPreludeにtoLowerがあるので、それを使って文字列に含まれる大文字をすべて小文字に直す関数を
tolowers :: String -> String tolowers s = map (toLower) s
と定義できる。
tolowers :: String -> String tolowers s = map (toLower) s solve :: String -> String -> Int solve a b = let (a',b') = (tolowers a,tolowers b) in if a' > b' then 1 else if a'==b' then 0 else -1 main = do a <- scan::IO String b <- scan::IO String putAnsLn (solve a b)
B:
中央のグリッドはどうしても通らなければならない＆それ以外ならいくらでも通らないことができる
ので、これも、中央のグリッドをマークされたマスが使ってしまっているかどうかだけ見ればいいです。
solve :: Int -> Int -> Int -> String solve n' x y = let n = div n' 2 in if (x==n||x==n+1)&&(y==n||y==n+1) then "NO" else "YES" main = do (n',x,y) <- scan::IO (Int,Int,Int) putAnsLn (solve n' x y)
C:
ある数yをn個に分割して、その平方の総和の値は
(y-n+1),1,1,1...,1(1がn-1個)
の時に最大となります。よって、これがxを超えるかどうか見ればそれでOKです。
厳密な証明はしてないです……正直通るかどうかは半分運でした。
rep 0 _ = [] rep i n = n:(rep (i-1) n) solve :: Integer -> Integer -> Integer -> [Integer] solve n x y = let max = y - (n-1) in if max < 1||(max*max+(n-1))<x then [-1] else max:(rep (n-1) (1::Integer)) main = do (n,x,y) <- scan::IO (Integer,Integer,Integer) putAnsLns (solve n x y)
このrepって関数、多分replicateで通ると思います。なんか書いてる途中で型が通らなくなって焦って代替の関数をかいてそのまま使ってしまいました。
なんやかんやで青復帰です。
このままDiv1に行けたらいいなーと思いますが、未だにDiv2のDが通らないのでまだ無理かもしれません。
それではまた。
Haskellでエラトステネスの篩(STUArray)
(9/1に追記)
(9/9に追記2)
そういや最近Haskell記事ばっかですね。
（まあでも、今回話題にする篩はC++でも以前記事にしましたね。）
STArrayを使う練習としてエラトステネスの篩を書いてみました。かなりすっきりと書けました。
速度に関する最適化は多少甘いかもしれない。
import Data.Array.ST import Data.Array.Unboxed import Data.Array.IArray import Monad sieve :: Int -> UArray Int Bool sieve n = runSTUArray $ do t <- newArray (0,n) True writeArray t 0 False writeArray t 1 False sequence_ [writeArray t i False|i<-[4,6..n]] let actions t k m = [readArray t i>>=(flip when) (writeArray t j False)|i<-[3,5..k],j<-[i*i,i*(i+2)..m]] let sqn = (floor.sqrt.fromIntegral) (n::Int) sequence_ (actions t sqn n) return t
これだけ。
「表から素数を取り出す→その素数の倍数をふるい落とす」の一連の操作をactionsという「モナド計算のリスト」にしてしまい、sequence_関数で一気に篩にかけています。これによりずいぶん見やすくなりました。
なお、sequence_関数（アンダースコア付き）を用いているのは結果がいらないからですが、sequence関数（アンダースコアなし）を用いるより数段早かったです。細かいところでも違いは出るものですね。
配列の解凍・凍結を繰り返すことなく一度のSTモナド内で計算するので、それなりに計算効率はいい……はず。
自分の環境では10000000要素の篩を作るのにGHC(-O2)でコンパイルすると0.240 secsほどでした。-O1オプションでも0.260 secでした。
なお、STArrayの配列更新がO(1)となるのは-O1,-O2オプションをつけてコンパイルしたときのみですので、GHCiや最適化オプション無しでは速度がガクっと落ちます。
オプション無しコンパイルでは10000000要素の篩を作るのに15秒程度かかってしまいました。
短いですが今回はこれだけで。
ではまた（・ω・）ﾉｼ
追記(9/1)
自分以前に篩を書いていた人の記事
「エラトステネスの篩」
http://d.hatena.ne.jp/mkotha/20101224
（www.kotha.netの裏）
を参考に、unsafeRead・unsafeWriteを用いて更に高速化。
これらの関数についてはほとんど資料がなかったのですが、Data.Array.Baseモジュールにて定義されている関数で、writeArray・readArrayとほぼ同じに使えますが、インデックスの指定はInt型でないと使えない、ということのようです。
readArray :: (MArray a e m, Ix i) => a i e -> i -> m e writeArray :: (MArray a e m, Ix i) => a i e -> i -> e -> m () unsafeRead :: (Ix i) => a i e -> Int -> m e unsafeWrite :: (Ix i) => a i e -> Int -> e -> m ()
また、sequence_よりmapM_を使った方がわずかに速かったです。
ということでコードは以下。
import Data.Array.ST import Data.Array.Unboxed import Data.Array.IArray import Monad import Data.Array.Base(unsafeRead, unsafeWrite) sieve :: Int -> UArray Int Bool sieve n = runSTUArray $ do t <- newArray (0,n) True unsafeWrite t 0 False unsafeWrite t 1 False mapM_ (\i -> unsafeWrite t i False) [4,6..n] let sqn = (floor.sqrt.fromIntegral) (n::Int) mapM_ (\i -> (mapM_ (\j -> unsafeRead t i>>=(flip when) (unsafeWrite t j False)) [i*i,i*(i+2)..n])) [3,5..sqn] return t
追記2(9/9)
unsafe操作以前に致命的な実装ミスをしていました……whenで条件判断するのは2つ目のmapM_の手前じゃないと、素数を発見して篩うときに毎回iが素数であるかのチェックをしてしまいますね。
そこを改善したコードが以下。
import Data.Array.ST import Data.Array.Unboxed import Data.Array.IArray import Monad import Data.Array.Base(unsafeRead, unsafeWrite) sieve :: Int -> UArray Int Bool sieve n = runSTUArray $ do t <- newArray (0,n) True unsafeWrite t 0 False unsafeWrite t 1 False mapM_ (\i -> unsafeWrite t i False) [4,6..n] 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
これなら1億篩うのに大体1sec程度で済みます。
Haskellにしては結構速いんではないでしょうか。
また高速化できたらここに追記します……
