Haskellでこの問題を解いてみた②

同じ人が別の問題を作っておられたようなので、解いてみました。

import Data.Array
import Data.List

type Point = (Int,Int)
type Field = Array Point Char
type Route = [Point]

isLegalPuyo :: Char -> Bool
isLegalPuyo '\n' = False
isLegalPuyo '\r' = False
isLegalPuyo _ = True

parseField :: String -> Field
parseField cs =
    let numRow = length $ lines $ cs
        numCol = length $ head $ lines $ cs
        puyos = filter isLegalPuyo cs
    in  array ((1,1),(numRow,numCol)) $ zip [(y,x) | y <- [1..numRow], x <- [1..numCol]] puyos

unfoldLine :: Int -> String -> [String]
unfoldLine x cs = case splitAt x cs of
                      (xs,[]) -> [xs]
                      (xs,ys) -> xs : unfoldLine x ys

showField :: Field -> IO ()
showField field = do
    putStr $ unlines $ unfoldLine (snd $ snd $ bounds field) $ elems field

dfs :: (a -> [a]) -> a -> [a]
dfs next node = dfs' node
    where
        dfs' n = n : concatMap dfs' (next n)

inBound :: Field -> Point -> Bool
inBound field (y,x) =
    let ((y1,x1),(y2,x2)) = bounds field
    in  y1 <= y && y <= y2 && x1 <= x && x <= x2

findNextRoutes :: Field -> Route -> [Route]
findNextRoutes field route =
    let pos@(y,x) = head route
        neighbors = filter (inBound field) [(y+1,x),(y-1,x),(y,x+1),(y,x-1)]
        puyo = field ! pos
        same_puyo_poss = filter (\p -> (field ! p) == puyo) neighbors
        not_dupe_poss = filter (\p -> not $ p `elem` route) same_puyo_poss
    in  map (:route) not_dupe_poss

findChainedRoute :: Field -> Point -> Route
findChainedRoute field pos = 
    let routes = dfs (findNextRoutes field) [pos]
    in  nub $ concat $ routes

erasePuyo :: Field -> Maybe Field
erasePuyo field =
    let puyo_poss = filter (\x -> (field ! x) /= ' ') (indices field)
        erase_routes = filter (\x -> length x >= 4) $ map (findChainedRoute field) puyo_poss
        erase_pos = nub $ concat $ erase_routes
    in  case erase_pos of
            [] -> Nothing
            xs -> Just (field // zip xs (repeat ' '))

dropPuyoLine :: Int -> Field -> Field
dropPuyoLine col field =
    let numRow = fst $ snd $ bounds field
        puyos = map (field !) [(p,col) | p <- [numRow,numRow-1 .. 1]]
        packedPuyos = filter (/= ' ') puyos
    in  field // zip [(p,col) | p <- [numRow,numRow-1 .. 1]] (packedPuyos ++ repeat ' ')

dropPuyo :: Field -> Field
dropPuyo field =
    let numCol = snd $ snd $ bounds field
    in  foldr dropPuyoLine field [1..numCol]

enumerateFields :: Field -> [Field]
enumerateFields field =
    case erasePuyo field of
        Just x -> let next_field = dropPuyo x
                  in  next_field : enumerateFields next_field
        Nothing -> []

showNumberedField :: (Int,Field) -> IO ()
showNumberedField (index,field) = do
    putStrLn $ show index
    putStrLn ""
    showField field
    putStrLn "------------------------------------"

main :: IO ()
main = do
    cs <- readFile "source.txt"
    let field = parseField cs
        enumerated_fields = enumerateFields field
    mapM_ showNumberedField (zip [1..] enumerated_fields)

だいたい1時間くらいでした。前の問題もそうだったんですが、慣れたＣ言語ならば半分くらいの時間で作れた気がするのが関数型言語好きとしては非常に悔しいです。関数脳を鍛えるには…。処理効率もあんまり良くないですし。

椿の日記

たぶんプログラムの話をします

Haskellでこの問題を解いてみた②