Haskell で AtCoder ABC444 A-E

はじめに

　実装メイン？　テンプレは別記事．

A

　問題文 : https://atcoder.jp/contests/abc444/tasks/abc444_a
　問題文をそのまま読むと文字列的に処理したい気持ちになりますが，入力が $3$ 桁で固定なので $111$ での剰余が $0$ かどうかで判定できます．ということで，入力の整数を答えに対応する真偽値に変換する関数は ( /= 0 ) . ( `mod` 111 ) です．この出力を対応する文字列にして出力すれば AC できます．

main = readInt >>= putStrLn . yesno . ( == 0 ) . ( `mod` 111 )

B

　問題文 : https://atcoder.jp/contests/abc444/tasks/abc444_b
　まず桁和ですが，（しばしば登場しますが）Data.Char にある digitToInt を使うと sum . map digitToInt . show という合成で実装できます．これを [ 1 .. n ] に map すると桁和からなるリストができるので，filter ( == k ) で条件を満たすものだけ取り出して length を取ると答えになります．

main = do
	[ n, k ] <- readInts
	print $ length $ filter ( == k ) $ map ( sum . map digitToInt . show ) [ 1 .. n ]

　おまけですが，桁和は unfoldr でも実装できます．unfoldr は fold された値からリストに戻すような関数で，

• リストの要素を取り出し終わったら（しばしば「単位元だったら」？）Nothing
• そうでないとき，リストに追加する値と次に持ち越す値のタプルを Just に入れたもの

を返す関数を渡して使います．今回の場合は $10$ での剰余をリストに追加して商を次に持ち越すので divMod が便利ですが，戻り値の要素の順序が逆なので swap します．$0$ には Nothing を返します．
　使い慣れてないのでエクササイズにと思ってやってみましたが，色々できそうな気配を感じます．再帰を unfoldr に任せられるのもよいですね．

digitSum = sum . unfoldr digitSum'
	where
	digitSum' 0 = Nothing
	digitSum' x = Just $ swap ( x `divMod` 10 )

C

　問題文 : https://atcoder.jp/contests/abc444/tasks/abc444_c
　最小値・最大値に着目したいのと，入力列の順序に意味がないので，$A$ は昇順ソート済みであるとします．
　$A_n$ が折れているか否かで場合分けして考えます．$A_n$ が折れているとき，$A_n$ の片割れは $A_1$ でなければならず，$l = A_1 + A_n$ です．両端を取り除いた列について再帰的に考えていくと$A$ と $A$ を逆順にしたもので綴じ合わせるしかなく，その結果として和が全て $l$ になっていればよさそうということになります．ちょっと慎重に考えると長さが偶数でないと同じもの同士を組み合わせてしまうので長さが偶数であることも必要条件です．
　$A_n$ が折れていない，すなわち $l = A_n$ の場合は $l$ と等しい要素には何もしなくてよいのでまず取り除きます．その後，残った列に対して $l = A_n$ として上記と同様のことをすれば判定できます．また，$l = A_1 + A_n$ のときに $l$ と等しい要素を取り除いても何も起きないので，仮定に拘らず $l$ と等しい要素を取り除いてしまえば実装を統一できます．
　ということで実装ですが，[ last as, head as + last as ] のそれぞれの要素について着目してそれを $l$ として，as' = dropWhileEnd ( == l ) as でできるリストを処理します．対応する要素同士を足し合わせる部分は zipWith (+) に as' と reverse as' を渡せばよく，長さのチェックは even とかで適当にやれます．全要素が $l$ と等しいかチェックする部分は map ( == l ) して and でもよいのですが，今回は練習として foldl してみます．畳み込みに使う関数は，第一引数にアキュムレータである Bool 値，第二引数に Int を受け取って次のアキュムレータを返すので，型は Bool -> Int -> Bool です．これを謎のこだわりでラムダを使わずに実装したいのですが，ちょっとテクいので先に答えを提示してしまいますと，( . ( l == ) ) . (&&) と実装できます．関数が Curry 化されていることを思い出すと過程を追うことができます：
\begin{align*}
( ( {} \circ ( l = {} ) ) \circ \mathord{ \land } )( a )( b )
&= ( ( {} \circ ( l = {} ) ) \circ ( a \land {} )( b ) \\
&= ( ( a \land {} ) \circ ( l = {} ) )( b ) \\
&= ( a \land ( l = b ) )
\end{align*}
ということで，あとは True を初期値として畳み込めばよいです．

main = do
	n <- readInt
	as <- sort <$> readInts
	printList do
		l <- [ last as, head as + last as ]
		let
			as' = dropWhileEnd ( == l ) as
			ss = zipWith (+) as' ( reverse as' )
		guard $ even ( length as' ) && foldl ( ( . ( l == ) ) . (&&) ) True ss
		return l

D

　問題文 : https://atcoder.jp/contests/abc444/tasks/abc444_d
　各桁に $1$ が何回足されたかを求めてから，下の桁から順に繰り上がりを処理します．
　桁毎に $1$ が足された回数を愚直に求めると TLE するので，この部分はいもす法をします．ということで，端点への足し引きはランダムアクセスができた方が楽なので accumArray で構築してからリストにして累積和をとります．accumArray に渡すリストは，各 a について [ ( 0, 1 ), ( a, -1 ) ] を集めてきて concat したものです．また，累積和は scanl1 (+) と簡単に実装できます．
　問題は繰り上がりの処理ですが，加算機みたいな気持ちになって繰り上がってきた値と次の桁を受け取って，次の繰り上がりとその桁の値を返す関数を考えると mapAccumL で畳み込めます．
　あとは，Leading Zero を落とすとか文字列に変換するとか細かい処理だけです．

main = do
	n <- readInt
	as <- readInts
	putStrLn $ map intToDigit $ dropWhile ( == 0 ) $ reverse $ snd $ mapAccumL carry 0 $ scanl1 (+) $ elems $ accumArray @UArray @Int (+) 0 ( 0, 300_000 ) $ concat do
		a <- as
		return [ ( 0, 1 ), ( a, -1 ) ]

carry c a = ( a + c ) `divMod` 10

E

　問題文 : https://atcoder.jp/contests/abc444/tasks/abc444_e
　着目している区間に含まれる要素の multiset をもってしゃくとり法をします．具体的には，区間の右側を伸ばしたいときに区間に入る値の内，それより大きいものの内最小なものと，それ以下のものの内最大のもののそれぞれとの差の絶対値が $d$ 以上なら伸ばせる，という判定をします．比較対象を二分探索でもってくることで TL に間に合います．
　ところが，二分探索できる multiset が標準っぽいところには無いので Data.IntMap で値と個数の対応を管理することで代用します．挿入の模倣の方は楽で，キーの存在性によって更新か挿入かを分岐してくれる関数 insertWith が使えます．他方，削除の模倣はちょっと面倒です．というのも，単に個数を $1$ 減らすだけだと $0$ 個になってもキーが残ってしまい，二分探索したときに multiset としては削除済みのものを発見することになって事故ります．消したい場合は消し，そうでない場合は更新をする関数 update があるので使えるのですが，場合分けはこちらで実装する必要があります．インターフェースとしては，値を受け取って，更新する場合は更新後の値を Just で包んだものを，キーを消す場合は Nothing を返す関数を渡します．今回の場合は，値が $1$ ならば消して，そうでないときは pred します．実装ですが，例によってラムダを避けて関数合成で構成していきます．真偽値を受け取って Just と Nothing に振り分ける部分には bool が使えます．ただし，単純に bool Nothing Just . ( /= 1 ) としてしまうと Nothing と Just で型が合いません．解決法としては Nothing の方を無理やり一引数関数にしてあげればよくて，使い方が謎と言われがち^[要出典]な関数 const を使ってbool ( const Nothing ) Just . ( /= 1 ) としてあげるとどっちも a -> Maybe a になって型が合います．個数をこの関数に渡して返ってくる関数に，個数を pred した値を渡してあげれば，update に渡すべき値になります．入力を上記の関数と pred のそれぞれに渡すのは Control.Arrow にある (&&&) が使えて，関数と引数のタプルにできます．タプル中の関数にタプル中の引数を適用するのは，($) の本来の意味を思い出せば uncurry ($) でできることが分かります．よって，全体としては uncurry ($) . ( bool ( const Nothing ) Just . ( /= 1 ) &&& pred ) となります．
　ということでデータ構造の準備ができたのでしゃくとり法を実装していきます．手続き型的に考えたときにアルゴリズム中で変化するものは，区間の右端点および左端点，区間内の要素からなる multiset（を模倣した Data.IntMap）です．この内，左端点は $1$ ずつ増えていくだけです．そこで残りの $2$ つだけを状態として考えると，次に着目する左端点を受け取って，右端点を伸ばせるだけ伸ばしたときの右端点とそのときの multiset を返す関数を考えることができます．そこから上述の方法で左端点を $1$ 進めれば新たな状態を作れるので，[ 1 .. n ] を mapAccumL で畳み込みながら各左端点毎の最右の右端点までの長さをリストにできます．

main = do
	[ n, d ] <- readInts
	as <- listArray @UArray ( 1, n ) <$> readInts
	let solve ( im, j ) i = ( ( IMap.update ( uncurry ($) . ( bool ( const Nothing ) Just . ( /= 1 ) &&& pred ) ) ( as ! i ) im', j' ), j' - i )
		where
		( im', j' ) = solve' ( im, j )
		solve' p@( im, j ) = if j <= n && le && gt
			then solve' ( IMap.insertWith (+) aj 1 im, succ j )
			else p
			where
			aj = as ! j
			le = let Just ( a, c ) = IMap.lookupLE aj im in d <= aj - a
			gt = let Just ( a, c ) = IMap.lookupGT aj im in d <= a - aj
	print $ sum $ snd $ mapAccumL solve ( IMap.fromList [ ( minBound `div` 2, 1 ), ( maxBound `div` 2, 1 ) ], 1 ) [ 1 .. n ]