Haskellでプログラミングコンテストチャレンジブック part1「くじびき」:二分探索

こんなタイトルなのはちょっと蟻本の勉強しようと思うから。主にHaskellC++ですね。

「くじびき」のコードをC++で書くのは容易いので今回はHaskellのみ。また、解いたのはサイズが大きいバージョンです。

問題概要:
整数値k_1……k_nから4つ選んで和をmにすることができるか。


入力:
n
m
k_1 k_2 k_3 ...


サイズ上限:
n:1000
m:10^8
k_i:10^8

Haskellで書くにあたってまず引っかかったのは二分探索。
C++ならほぼ迷いなく書けるのだけど、っていうかSTL(もっといえばCのライブラリ関数)にあるんだけど、Haskellの配列の扱いがなんかややこしくて手間取りました。

ライブラリ化してすぐ使えるようにしておくのが良さそうですね。

以下コード
nは捨てました。forもないし使わないので。

import Array
import List

sortArray :: (Ord a) => [a] -> Array Int a
sortArray l = let leng = length l - 1 in array (0,leng) (zip [0..leng] (sort l))

binary_search :: (Ord a) => Array Int a -> a -> Bool
binary_search ary e = bsearch begin end
    where
      (begin, end) = bounds ary
      bsearch n m = 
          if m < n + 1 then
              False
          else
              let mid = ((m+n) `div` 2)
              in
                let midval = (ary!mid)
                in
                  if midval > e then
                      bsearch n mid
                  else if midval < e then
                           bsearch (mid+1) m 
                       else 
                           True

solve n l = let sums = sortArray [i+j|i<-l,j<-l]
            in
              if (or [binary_search sums (n-i-j) | i<-l,j<-l]) then
                  "Yes"
              else
                  "No"

convToInts :: String -> [Int]
convToInts l = [read n::Int|n <- (words l)]

main = do getLine
          m <- getLine
          line <- getLine
          putStrLn (solve ((read m)::Int) (convToInts line))


sortArraybinary_searchが今回のメイン。

sortArrayはリストからソート済み配列を得る関数。

binary_searchはソート済み配列aryに要素nが含まれているかどうかを判定する関数。

関数arrayは範囲(begin,end)と配列の内容listを引数にとって配列を返す関数。前にLCSをHaskellで書いてみた時も使いました。

配列の内容は連想リスト[(i,k_i)](i:インデックス,k_i:値)で渡しますが、このインデックスにはIxのインスタンスなら何でも使えるそうです。まあ、現実的には整数か整数のタプルを使うことになるのでしょうが。

binary_search関数はCやC++での実装とほぼ同じなはずです。Haskellのarrayのアクセス時間がO(1)なら、これで計算量O(log n)が達成できているはずです。


まあこれで実装はできますが、実行時間がちょっと遅いです。
計算量はO(n^2 log n)を達成できているはずなのですが、n=1000で"No"を出力するまでに大体5秒くらい。
これじゃあ、時間制限によっては落ちますね。


今回はここまで。
はたしてpartいくつまで続くかな。

それでは。