Haskellでプログラミングコンテスト チャレンジブック part２:「Lake Counting」

Haskellで蟻本記事第二弾。

Haskellでプロコンの問題に取り組もうと思ったときにまず難しいと思ったのは、基本的にリストでデータの集合を扱うHaskellでは二次元データへのO(1)でのアクセスをやろうとすると面倒なこと。
次に、そういったデータの集合を更新しながら進むアルゴリズムがやるづらいだろうということ。破壊的代入ができないので。

と、いうことでこの問題"Lake Counting"。
POJの問題で問題Noは2386。

「庭」の大きさを表す整数N,M(N:縦、M:横)が一行で与えられ、続けて以下のような「庭」のデータが与えられる。「w」は「水たまり」を表し、周囲８マスで隣接した水たまりはつながっているとみなす。

...w..............wwww........

..www............w...w........

..www.............w..w........

...w.............w...w........

................w....w........

.....wwwwww......w...w........

....ww............wwww........

...w..........................

..w..wwwwwwwwwww.....ww.......

.w...wwwwwwwwwww....wwww.....w

w....wwwwwwwwwww.....ww.....ww

w....wwwwwwwwwww..........wwww

w........................wwwww

w.....w.w.w.w..........wwwwwww

w....................wwwwwwwww

水たまりがいくつあるかを出力せよ。（上の例では10）

データサイズ
N,M≦100

解法は、データを頭から見ていって「w」に行き着いたら深さ優先探索でそこに隣接する「w」をすべて「庭」から取り除いていき、先頭のループで何回DFSを呼び出したかをカウントするだけ。

手続き型言語なら二次元データ＆配列の破壊的更新でDFSを書くことになりますが、Haskellでも同じようなことができるか実験。

Haskellによるコードは以下。

import Array

getLines :: Int -> IO [String]
getLines n = if n > 0 then
                 do l <- getLine
                    ls <- getLines (n-1)
                    return (l:ls)
             else
                 return []

getIntLine :: IO Int
getIntLine = do n <- getLine
                return (read n::Int)

putAnsLn :: (Show a) => a -> IO ()
putAnsLn ans = putStrLn (show ans)

matrix :: Int -> Int -> [[a]] -> Array (Int,Int) a
matrix n m lis = array ((0,0),(n-1,m-1)) (matrix_rec 0 lis [])
    where
      matrix_rec _ [] l = l
      matrix_rec i (x:xs) l = matrix_rec (i+1) xs ((zip [(i,j)|j<-[0..m]] x)++l)

{- DFS -}
dfs :: Int -> Int -> Int -> Int -> Array (Int,Int) Char -> Array (Int,Int) Char
dfs x y n m garden_init = dfs_rec (garden_init//[((x,y),'.')]) [(nx,ny)|(nx,ny) <- [(x-1,y-1),(x,y-1),(x+1,y-1),(x-1,y),(x+1,y),(x-1,y+1),(x,y+1),(x+1,y+1)] , nx>=0 , nx<n , ny>=0 , ny<m]
    where
      dfs_rec garden [] = garden
      dfs_rec garden ((nx,ny):cs) = if garden!(nx,ny) == 'w' then
                                        let new_garden = dfs nx ny n m garden
                                        in
                                          dfs_rec new_garden cs
                                    else
                                        dfs_rec garden cs

solve :: Int -> Int -> [String] -> Int
solve n m gardenlist = let garden = matrix n m gardenlist {- 入力を二次元配列に直す -}
                       in
                         solve_rec 0 0 garden 0
                             where
                               solve_rec i j gard ans = if i<n then
                                                            if j<m then
                                                                if gard!(i,j) == 'w' then
                                                                    {- 水たまりが見つかったらDFSでその水たまりを庭から取り除く -}
                                                                    solve_rec i (j+1) (dfs i j n m gard) (ans+1)
                                                                else
                                                                    solve_rec i (j+1) gard ans
                                                            else
                                                                solve_rec (i+1) 0 gard ans
                                                        else
                                                            ans

main = do n <- getIntLine
          m <- getIntLine
          garden <- getLines n
          putAnsLn (solve n m garden)

演算子「//」は「ary//[(i,new_data)]」で「配列aryのi番目の要素にnew_dataを代入した配列を返す」演算子で、配列の漸進更新ができる。これによって破壊的代入をするようなコードも書ける。（実際の挙動はどうなっているのかわからないが）
（追記：どうやら、破壊的代入ではなく配列全体を作り直すような実装になっており、配列の漸進的更新のコストはO(N)となっているようです。あとで改良したコードを載せますが、上のコードの計算量はかなり大きくなってしまっています。）

実際書いてみて思ったのは、やはり二次元配列を作るときですら多少面倒な手順を踏まなくてはならずデータを使用可能な形に落とし込みづらかったのと、２重ループが書きづらかったということ。

getLines 〜 matrixはライブラリ化しておいたほうがいいですね。特にIO周りで手こずるわけには行きませんので。

今夜Codeforcesがありますね。参加しようかなぁ。

それではﾉｼ