ESM オフラインリアルタイムどう書くの問題を解いてみた（Haskell版）

19475 ワード
どう書く Haskell Haskell テキストリンク

はてブで見かけたので Haskell で解いてみました。

問題を素直にコードにすることを意識してソースの通り、ほぼ、上から順にテストしながら書き下して行きました。最終的な動作確認までの所要時間は40分位です。

Haskell の復習を兼ねていたので良い題材になりました。

せっかくなので、実装の思考過程を書いておきます。

データ構造と関数の型を考える

最初にラーメン屋、席、お客さんの列のデータ型から考えました。

ルールにある通り、席の状態が遷移するので素直に席（Chair）の型はそれぞれの状態を列挙型になるようにしました。

data Chair = Waiting | Eating | Resting | Empty
    deriving (Enum, Eq, Ord, Bounded, Show)

状態をすすめるとき、単純に succ 関数を適用すれば次の状態になるわけですが、 Empty から次の状態になるには客が席につくときに限るので、succ　関数で次の状態にならないように最後に定義しました。（succ Empty とすると例外が発生します）

ラーメン屋（Store）は席のリストとし、お客さんの列（Customers）は人数のリストにしました。

type Store = [Chair]
type Customers = [Int]

次に主要な関数をざっくり考えました。まずは席の状態をすすめる関数。これはChair -> Chairになり、上記の通り succ関数ですすめる。空席（Empty）だけは次の状態も空席のままにするだけの簡単な関数なのですぐに実装。

-- 席の状態をすすめる
stepForward :: Chair -> Chair

次に、店の状態を新しくしてお客を入れていく関数。客の行列が空になるまで再帰して最後の店の状態が今回の答えになります。
引数は店の状態と客の行列、戻り値として新しい店の状態を返すので、関数の型は Store -> Customers -> Store となる。
新しい店の状態は、引数の店（Store）の各席に stepForward 関数を適用すればいい。

-- 行列がなくなるまで店の状態をすすめる
simulate :: Store -> Customers -> Store

客を入れるには、客の数分、連続して空いている席があるか探す必要があるけど、実装がぱっと思い浮かばなかったので関数を切り出した。
引数に必要な席数と店の状態、戻り値は席の位置を返すようにしてみた。見つからない場合もあるので戻り値の型は Maybe Intにした。 

-- 空き席を探す
searchSpace :: Int -> Store -> Maybe Int

searchSpace 関数の実装

simulate は新しい店の状態を計算して、次の先頭の客が座れるときは行列の先頭を取り除いて再帰、そうでなければ行列をそのままで再起すれば次の状態になるので、悩むことはない。

searchSpace の実装が一番悩みました。方針としては、必要な数だけ Empty が並んでいるかスキャンする方法を考えてみた。
初めに思いついたのは foldl を使って、ひとつ前の値と連続している数のデータを持ち回して畳み込むことを考えたが、煩雑な記述になりそうだったの却下。
Data.List モジュールを眺めると、連続した要素をまとめたリストを作る group 関数が見つかったのでこれを使ってみました。

group [Eating, Empty, Empty, Waiting, Waiting]     -- [Chair]
  = [[Eating], [Empty, Empty], [Waiting, Waiting]] -- [[Chair]]

group の結果から [(Eating, 1), (Empty, 2), ...] みたいに値と連続した要素の数の組を作って、そのリストをスキャンすれば、必要な空き席が連続しているか見つけられそう。ずばりfindIndex 関数で条件にあった位置が検索できる。

findIndex :: (a -> Bool) -> [a] -> Maybe Int

findIndex で見つけた位置は、group 関数が返したリストの位置なので、その位置の前までの要素の2番めの値を集計すれば、元のリストの位置に変換できておしまい。

sit 関数の実装とバグフィックス

席の空きが確認できたら、店の状態を人数分 Empty から Waiting に状態を更新しなければならない。最初は、 simulate 関数の中で定義してたけど、テストしにくいので sit 関数として抽出した。
これは、先に検索した位置を元に店の状態を切り貼りして新しい状態のリストにすればいいので簡単。

ここで一通り完成したので、テストデータを使って動作確認をしたところ、いくつか失敗しました。よくよく確認すると、席は円になってて店の状態は循環したリストにしないといけないことに気づいた。

単純に、searchSpace 関数で店の状態を store ++ store のように繋いで位置を検索するように修正。あわせて sit 関数で切り貼りするときに、後ろと前が繋がるような切り貼りにするように修正して対応。

席の空きを検索するのと、客が座った状態にするのは一緒にしたほうが効率がいいような気がしたけど、実装のスピードを優先した。

完成したソースコード

module Main where

import Data.List
import Data.Char

data Chair = Waiting | Eating | Resting | Empty
    deriving (Enum, Eq, Ord, Bounded, Show)

type Store = [Chair]
type Customers = [Int]

-- 席の状態をすすめる
stepForward :: Chair -> Chair
stepForward c
 | maxBound == c = c
 | otherwise = succ c

-- 空き席を探す
searchSpace :: Int -> Store -> Maybe Int
searchSpace 0 _ = Nothing
searchSpace n s =
    fmap (\x -> foldl (\b a -> b + snd a) 0 $ take x runlength) index
    where
        runlength = zip (map head ys) (map length ys)
        ys = group $ s ++ s
        index = findIndex (\x -> Empty == fst x && n <= snd x) runlength


-- 行列がなくなるまで店の状態をすすめる
simulate :: Store -> Customers -> Store
simulate store [] = store
simulate store cs@(h:t) =
    case searchSpace h next of
        Just index -> simulate (sit index h next) t
        Nothing -> simulate next cs
    where
        next = map stepForward store

-- 客を座らせる
sit :: Int -> Int -> Store -> Store
sit i numOfCustomers store =
    replicate firstHalf Waiting
    ++ (take (i - firstHalf) $ drop firstHalf store)
    ++ replicate secondHalf Waiting
    ++ drop (i + secondHalf) store
    where
        firstHalf = if i + numOfCustomers > length store then numOfCustomers - (length store - i) else 0
        secondHalf = numOfCustomers - firstHalf

-- 以下、テスト用
storeState :: Store -> String
storeState = mconcat . map (\c -> if c == Empty then "0" else "1")

toCustomers :: String -> Customers
toCustomers = map digitToInt

test :: String -> String -> IO ()
test input expected =
    if expected == actual then
        putStrLn "Ok"
    else
        putStrLn $ "Fail: input=" ++ input ++ " expected=" ++ expected ++ " actual=" ++ actual
    where
        actual = storeState $ simulate (replicate 8 Empty) $ toCustomers input

main = do
    test "3316" "11111110"
    test "3342153" "11111111"
    test "8" "11111111"
    test "232624" "11110110"
    test "1" "10000000"
    test "224673432" "11111000"
    test "123464334" "11111110"
    test "44372332" "11111111"
    test "2344" "11111111"
    test "6448476233" "11111100"
    test "4345174644" "11111111"
    test "77743" "11111110"
    test "2136524281" "10000000"
    test "344373" "11100000"
    test "434226" "11111111"
    test "7433223" "11111110"
    test "2246534" "11110111"
    test "634" "11111110"
    test "2222" "11111100"
    test "2442343238" "11111111"
    test "7243344" "11111111"
    test "26147234" "10111111"
    test "34424" "10011111"
    test "6334" "11011111"
    test "3828342" "11011110"
    test "4431" "11110000"
    test "2843212125" "11111111"
    test "333336482" "11000000"
    test "374" "11110000"
    test "4382333" "11100111"
python出力カラーテキスト
pickleとcPickle-pythonオブジェクトのシーケンス化