Haskell で AtCoder ABC447 A-F

はじめに

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

A

　問題文 : https://atcoder.jp/contests/abc447/tasks/abc447_a
　端から詰めていくと，最後の一人を除いて $2$ 席を占有し最後の一人は $1$ 席だけ使うので，必要な座席は $2m - 1$ 席です．よって，$2m - 1 \leq n$ かどうかで判定できます．

main = do
	[ n, m ] <- readInts
	putStrLn $ yesno $ 2 * m - 1 <= n

B

　問題文 : https://atcoder.jp/contests/abc447/tasks/abc447_b
　文字毎に個数を数えてから，個数が最大値と異なるものを filter します．文字から個数への索引は添字として文字を直接渡せた方が楽なので Array にします．ミュータブルの方でよいので accumArray で作ることにすると，文字が現れる毎に $1$ を加算すればよいので accumArray (+) 0 ( 'a', 'z' ) に zip s $ repeat 1 を渡せばよいです．

main = do
	s <- getLine
	let
		cs = accumArray @UArray @Int (+) 0 ( 'a', 'z' ) $ zip s ( repeat 1 )
		ma = maximum $ elems cs
	putStrLn $ filter ( ( /= ma ) . ( cs ! ) ) s

C

　問題文 : https://atcoder.jp/contests/abc447/tasks/abc447_c
　それぞれの先頭同士から始めて $1$ 文字ずつマッチさせます．操作によって読み飛ばせるので，

• 文字が一致しているならコスト $0$ でマッチさせる
• 片方が A ならコスト $1$ で A の方だけ読み飛ばす
• どちらもでない（両方とも A でなく，一致もしていない）とき，不可能

となります．片方のみ最後まで読んだ場合は，

• 他方の文字が全て A ならその長さ分のコストで読み飛ばせる
• そうでないなら不可能

です．
　以上のことを再帰で処理しますが，失敗するかもしれない計算ということで Maybe を使うことにします．Maybe Int を返すように実装してもよいのですが，今回は [ Maybe Int ] を返してから加工します．sequence を Maybe のリストに対して使うと，すべてが Just 値なら中身を並べたリストを Just に包んだものを，そうでなければ Nothing が帰ってきます．これに fmap sum すると不可能なら Nothing が，そうでなければ答えが入った Just 値が得られます．あとは fromMaybe -1 に通してあげれば出力するべきものになります．

main = do
	s <- getLine
	t <- getLine
	print $ fromMaybe -1 $ fmap sum $ sequence $ solve s t

solve [] [] = [ Just 0 ]
solve s [] = [ lengthA s ]
solve [] s = [ lengthA s ]
solve s1@(c1:cs1) s2@(c2:cs2)
	| c1 == c2  = solve cs1 cs2
	| c1 == 'A' = Just 1 : solve cs1 s2
	| c2 == 'A' = Just 1 : solve s1 cs2
	| otherwise = [ Nothing ]

lengthA s = if all ( == 'A' ) s then Just $ length s else Nothing

D

　問題文 : https://atcoder.jp/contests/abc447/tasks/abc447_d
　公式解説 にある方法を実装します．実装面ではほとんど言うことが無くて，トリプルと次の文字を受け取って次ステップのトリプルを返す関数を作って畳み込めばよいです．
　微妙な非自明ポイントとしては，トリプルの第 $3$ 要素を取り出す関数 thd3 は Data.Tuple.Extra にあることでしょうか．

main = getLine >>= print . thd3 . foldl solve ( 0, 0, 0 )

solve ( a, b, c ) x = case x of
	'A' -> ( succ a, b, c )
	'B' -> ( a, min a $ b + 1, c )
	'C' -> ( a, b, min b $ c + 1 )

E

　問題文 : https://atcoder.jp/contests/abc447/tasks/abc447_e
　ある辺 $( u, v )$ を採用しても連結成分が $1$ 個にならない条件は，

• その時点での連結成分数が $3$ 以上
• $u, v$ が同じ連結成分に属する

のいずれかを満たすことです．辺重みの降順にこの条件を判定して，採用できるものはすべて採用することで最適解になります．
　辺を追加しつつ連結性の判定をするために Disjoint Set Union（いわゆる Union-Find）が欲しいので，ac-library-hs にある AtCoder.Dsu を使います．AtCoder.Dsu は PrimMonad の中で動くのと，連結成分の個数の取得は groups 経由でやると遅い（たぶん）ので別に管理する必要があるため，ST で実装します．
　連結成分数が変化するのは異なる連結成分に属する頂点間を結ぶ辺を追加したときですが，mergeMaybe でマージすると戻り値が Nothing か否かによって判定できます．
　あとは実装ですが，モナドアクションの中身を使って条件分岐するには Control.Monad.Extra にある ifM が便利で，$2$ つの条件の OR を取る部分は <$>, <*> を使うと中間変数無しに書けます．

import qualified AtCoder.ModInt as MINT
import qualified AtCoder.Dsu as DSU

-- テンプレ中略

main = do
	[ n, m ] <- readInts
	es <- map ( map pred ) <$> replicateM m readInts
	print $ sum $ solve n es

solve n es = runST do
	dsu <- DSU.new n
	g <- newSTRef n
	forM ( reverse $ zip [ 1 .. ] es ) $ \( i, [ u, v ] ) -> do
		ifM ( (||) <$> ( ( 2 < ) <$> readSTRef g ) <*> DSU.same dsu u v )
			( do
				modifySTRef g =<< subtract . ( bool 0 1 . isJust ) <$> DSU.mergeMaybe dsu u v
				return 0
			) (
				return ( MINT.pow 2 i :: MINT.ModInt998244353 )
			)

F

　問題文 : https://atcoder.jp/contests/abc447/tasks/abc447_f
　適当な頂点を根にして根から DFS をして，各頂点で

• その頂点以下の部分木で完結しているムカデグラフの最大サイズ
• その頂点を経由して根方向に伸びることができるムカデグラフの最大サイズ

の $2$ つを計算します．根における前者の値が答えです．
　DFS をする都合グラフは隣接リストで受け取りたいですが，辺を双方向化してから accumArray ( flip (:) ) [] で配列にすると実装できます（(:) が可換ではないので flip が要ります）．DFS は根方向への逆走を防ぐために直前の頂点（根のときは無いので適当な値）p と今の頂点 u の両方を受け取るようにしますが，このとき，子たちへの再帰呼び出しは map を使って map ( dfs u ) $ filter ( /= p ) ( graph ! u ) と実装できます．

main = readInt >>= flip replicateM do
	n <- readInt
	graph <- accumArray @Array ( flip (:) ) [] ( 1, n ) . concat <$> replicateM ( n - 1 ) do
		[ u, v ] <- readInts
		return [ ( u, v ), ( v, u )]
	print $ solve graph

solve graph = fst $ dfs -1 1
	where
	n = snd $ bounds graph
	dfs p u = ( res1, res2 )
		where
		d = length $ graph ! u
		( dp1, dp2 ) = ( id *** reverse . sort ) $ unzip $ map ( dfs u ) $ filter ( /= p ) ( graph ! u )
		res1 = max ( maximum $ 0 : dp1 ) if
			| 4 <= d -> max 1 $ dp2 !! 0 + dp2 !! 1 + 1
			| d == 3 -> max 1 $ dp2 !! 0 + 1
			| d == 2 -> 1
			| otherwise -> 0
		res2 = if
			| 4 <= d -> dp2 !! 0 + 1
			| d == 3 -> 1
			| otherwise -> 0