Haskell で AtCoder ABC 401 A-D

はじめに

　テンプレは別記事．

A

　問題文 : https://atcoder.jp/contests/abc401/tasks/abc401_a
　$s \in [ 200, 299 ] \cap \mathbb Z$ か否かを何らかの方法で判定するだけです．範囲が狭いので `elem` [ 200 .. 299 ] で済ませました．

main = readInt >>= putStrLn . bool "Failure" "Success" . ( `elem` [ 200 .. 299 ] )

B

　問題文 : https://atcoder.jp/contests/abc401/tasks/abc401_b
　ログイン状態と操作を受け取って，操作後のログイン状態と今回の操作で発生する認証エラーの回数 ($0$ or $1$) を返す関数を作ると mapAccumL で畳み込めます．
　ログイン状態を $0$，非ログイン状態を $1$ で表すことにすると記述量が減りますが，セマンティクス的に微妙かもしれません……？（最初は真面目に Bool にしていた）

main = readInt >>= flip replicateM readStr >>= print . sum . snd . mapAccumL solve 1

solve _ "login" = ( 0, 0 )
solve _ "logout" = ( 1, 0 )
solve l "public" = ( l, 0 )
solve l "private" = ( l, l )

C

　問題文 : https://atcoder.jp/contests/abc401/tasks/abc401_b
　Fibonacci 数列と言えば Haskell では以下の有名コードがありますね：

fs = 1 : 1 : zipWith (-) ( tail fs ) fs

遅延評価がうまいことやってくれて先頭から順に値が決まっていくというやつです．今回の問題でも似たようなことはできないでしょうか……？
　今回の場合は $k$ 項に渡る和を（制約的に）$O( 1 )$ 時間で（なくてもいいけど十分高速に）求められるようにする必要があるので，累積和を使います．累積和は，Haskell 的に言えばリスト as に対して psum = scanl (+) 0 as というリストです*1．このとき，各 $i$ に対して $i$ 項目から $i + k - 1$ 項目までの $k$ 項の和からなるリストは zipWith (-) ( drop k psum ) psum になります．累積和を表すリストの連続する $k + 1$ 項が決まっていればこのような zipWith ができますが，$A$ の先頭で連続する $k$ 項の $1$ と scanl の初期値の $0$ があるので足りています．なので，作りたいリストの先頭部分を replicate k 1 として，その後ろに zipWith (-) ( drop k psum ) psum を連結すれば欲しい列がうまいこと（ほぼ）求まります．今回の問題は ${} \bmod { 10^9 }$ を求める必要があるので，演算子を定義するなどする必要はあります．

main = do
	[ n, k ] <- readInts
	print $ solve k !! n

solve k = as
	where
	as = replicate k 1 ++ zipWith (-%) ( drop k psum ) psum
	psum = scanl (+%) 0 as

a +% b = ( a + b ) `mod` 10^9
a -% b = ( a - b + 10^9 ) `mod` 10^9

D

　問題文 : https://atcoder.jp/contests/abc401/tasks/abc401_d
　まず与えれた文字列 $S$ において o に隣接する ? は . にしかできないので，まずこの前処理をします．リストでうまいことやってもよいのですが，先頭・末尾の処理が面倒そうなので $S$ を listArray してしまって添字アクセスでやることにします．各位置 $i$ について隣の要素が存在し，かつ o かどうかを判定する必要がありますが，Data.Array.IArray に範囲チェック付きの添字アクセスをする演算子 (!?) がある（結果は Maybe で帰ってくる）があるので s !? ( i - 1 ) == Just 'o' のように書けていい感じ……と思いきや AtCoder 側の GHC が古いか何かの理由で存在しないので，以下のように自作しました（テンプレに入れた）：

a !? i
	| inRange ( bounds a ) i = Just $ a ! i
	| otherwise = Nothing

　前処理が終わったあと，連続する ? が一意的に定まるケースが 2 つあります：

• 入力に $k$ 個の o が含まれる．
• 詰められる o の最大数だけ o を追加する必要がある．

前者は簡単で，全ての ? を . にします．後者の場合は，まず詰められる最大数がいくつになるかというと，? が連続する極大な部分の長さが $l$ のとき，$\frac { l + 1 } 2$ 個の o を入れられるので，和を取ると入れられる最大数が分かります．最大限詰め込むとき，$l$ が奇数ならばならば入れ方は一意です（そうでないなら，奇数文字目に入れるか偶数文字目に入れるかの $2$ 通りがある）．連続する同じ文字からなる極大な部分列に分解する部分は group が便利です．あとは……がんばります．

main = do
	[ n, k ] <- readInts
	s <- listArray ( 1, n ) <$> readStr
	let
		s' = do
			i <- [ 1 .. n ]
			return $ if s !? ( i - 1 ) == Just 'o' || s !? ( i + 1 ) == Just 'o'
				then '.'
				else s ! i
		k' = k - count 'o' s'
		gs = map ( head &&& length ) $ group s'
		filled = sum $ map ( ( `div` 2 ) . succ ) $ map snd $ filter ( ( == '?' ) . fst ) gs
	putStrLn $ if k' == 0
		then replace "?" "." s'
		else concatMap ( solve $ filled == k' ) gs

solve f ( c, l )
	| not f || c /= '?' = replicate l c
	| odd l = take l $ cycle "o."
	| otherwise = replicate l '?'