これは何?
上記アドベントカレンダーの11日目の記事です。
この記事では、Haskellで入力処理を高速化するテクニックについて紹介します。
interactでIO処理をまとめる
haskell-jpのもくもく会で教えていただいたYorgey先生のブログで紹介されているテクニック。
interactの型はinteract :: (String -> String) -> IO () であり、入力を受けとり、出力を返す処理を書くと勝手にIO処理をやってくれるという便利な関数。
とりあえず、以下の問題を解いた際のコードを例として貼っておく。
import Data.List ((\\))
toTuple :: [String] -> (Char, Int)
toTuple [suit, rank] = (head suit, read rank)
solve :: [(Char, Int)] -> String
solve cards = unlines $ map (\(s, n) -> s : ' ' : show n) $ findMissingCard cards
findMissingCard :: [(Char, Int)] -> [(Char, Int)]
findMissingCard cards =
let fullDeck =
[ (s, n)
| s <- "SHCD",
n <- [1 .. 13]
]
in fullDeck \\ cards -- \\はリストの差を求める演算子
main :: IO ()
main = interact $ \inputs ->
let ls = lines inputs
_ = read $ head ls :: Int
cards = map (toTuple . words) $ tail ls
in solve cards
interactは改行を含むことに注意
main :: IO()
main = interact $
show
こういった簡単なサンプルを使って試してみるとわかるのだが、getLine等と異なり、interactは改行を含んでいる。
AtCoder\n
ストリーミングが狙えると速く出来る
Note that this operation is lazy, which allows to produce output even before all input has been consumed.1
interactは遅延評価のため、全てのinputを消費する前にoutputを出力することが可能である。
このため、ストリーミングを意識してプログラミングを行うことができると、高速にできる。
ストリーミングが使えた問題を例に、ストリーミングで高速化できた例を紹介する。
-- ChatGPTに作ってもらったストリーミングを使わないシンプルな実装
import Data.Char (isLower)
solve :: String -> String
solve s
| head s /= 'A' = "WA"
| otherwise =
let mid = drop 2 (take (length s - 1) s) -- 3文字目〜末尾-1
in case filter (== 'C') mid of
[_] | all isLower (filter (/= 'C') (tail s)) -> "AC"
_ -> "WA"
main :: IO ()
main = do
s <- getLine
putStrLn (solve s)
-- interact使用バージョン
import Control.Arrow ((>>>))
import Data.Char (isUpper)
solve :: String -> String
solve (s0 : s1 : rest)
| s0 /= 'A' = "WA" -- 先頭がAでない
| isUpper s1 = "WA" -- 2文字目は大文字でない
| otherwise = go rest False
where
-- 文字列 先頭から3文字目から末尾の2文字目までの間にcがちょうど1つあるか
-- NOTE: interactは改行文字列\nを含むため、改行手前まで読めば終了する
go ['\n'] isC = if isC then "AC" else "WA"
go (c : rest) isC
| not isC && c == 'C' =
case rest of
['\n'] -> "WA" -- 末尾にCがあるのはNG
_ -> go rest True
| isC && c == 'C' = "WA" -- 大文字のCが2回以上登場するのはNG
| isUpper c && c /= 'C' = "WA"
| otherwise = go rest isC
main :: IO ()
main =
interact $
solve >>> (++ "\n")
比較するとストリーミングを使った実装のほうが速い。
この問題をストリーミングを使った実装にするに当たり、ロジックの記述が複雑になり、ACするまでにかなりの時間を要してしまった。
競技プログラミングの場合、速く動作するだけでなく、速くACすることも重要なので必ずしも2つ目の実装を推奨するわけではないことに留意されたし。
補足: ストリーミングとは関係なく、IO処理が一箇所にまとまったことで高速化された例
実際のコンテストの問題で比較してみる。
この問題では入力に対してsortを行っているため、sortを実施するためにすべての入力が終わるのを待つ必要がある。
そのため、interactを使ってもストリーミングによる効率化は難しい。
だが、このようなケースでもinteractを使うと多少早くなることがあったので参考までに記載しておく。
-- interactバージョン
import Data.List (sort)
check :: [Int] -> [Int] -> Int -> Bool
check hs bs k = go hs bs 0
where
go [] _ cnt = cnt >= k
go _ [] cnt = cnt >= k
go (h:hs') (b:bs') cnt
| h <= b = if cnt + 1 == k
then True
else go hs' bs' (cnt + 1)
| otherwise = go (h:hs') bs' cnt
main :: IO ()
main = interact $ \input ->
let xs = map read (words input) :: [Int]
n = xs !! 0
m = xs !! 1
k = xs !! 2
hs = sort $ take n $ drop 3 xs
bs = sort $ take m $ drop (3 + n) xs
in unlines [if check hs bs k then "Yes" else "No"]
-- interactなしバージョン
import Data.List (sort)
check :: [Int] -> [Int] -> Int -> Bool
check hs bs k = go hs bs 0
where
go [] _ cnt = cnt >= k
go _ [] cnt = cnt >= k
go (h:hs') (b:bs') cnt
| h <= b = if cnt + 1 == k
then True
else go hs' bs' (cnt + 1)
| otherwise = go (h:hs') bs' cnt
main :: IO ()
main = do
[n, m, k] <- map read . words <$> getLine
hs <- sort . map read . words <$> getLine
bs <- sort . map read . words <$> getLine
putStrLn $ if check hs bs k then "Yes" else "No"
メモリ使用量および、実行時間が改善されているのが確認できる。
interactがないバージョンだと何度もgetLineを呼び出しているのが遅くなっている原因と考えられる。
もちろんgetContent等を使えばgetLineを何度も呼び出す必要はないのだが、個人的にIO処理を一箇所にまとまるとコードの見通しが良くなるので気に入っている。
Reference

