はじめに
実装メイン．テンプレは別記事．
A
問題文 : https://atcoder.jp/contests/abc441/tasks/abc441_a
Yes になる必要十分条件は，$p \leq x \land x < p + 100 \land q \leq y \land y < q + 100$ です．そのまま実装してもよいですが，inRange を使うと意図が明確になります（微差な気もしますが）．
main = do [ p, q ] <- readInts [ x, y ] <- readInts putStrLn $ yesno $ inRange ( ( p, q ), ( p + 99, q + 99 ) ) ( x, y )
B
問題文 : https://atcoder.jp/contests/abc441/tasks/abc441_b
単語 $W$ についての出力は，
\begin{align*}
\phi &= \text{$W$ の全文字が $S$ に含まれる} \\
\psi &= \text{$W$ の全文字が $T$ に含まれる}
\end{align*}
とすれば
- $\phi \land \lnot \psi \rightarrow$
Takahashi - $\lnot \phi \land \psi \rightarrow$
Aoki - $\text{otherwise} \rightarrow$
Unknown
となります．
文字列の全要素が s に含まれるかを判定する関数は all ( `elem` s ) で実装できるので，あとは条件分岐を書けば問題を解けます．
main = do [ n, m ] <- readInts s <- readStr t <- readStr q <- readInt ws <- replicateM q readStr mapM_ putStrLn do w <- ws let ft = all ( `elem` s ) w fa = all ( `elem` t ) w return if | ft && not fa -> "Takahashi" | not ft && fa -> "Aoki" | otherwise -> "Unknown"
C
問題文 : https://atcoder.jp/contests/abc441/tasks/abc441_c
最大限意地悪をされた場合，量が多い方から $n - k$ 個が水でそれ以外がぽしゃけになります．また，確実にぽしゃけである部分も量が多い方から飲むのが最適です．
なので，$A$ を降順ソートしてから先頭 $n - k$ 要素を取り除き，残りの要素の累積和をとって初めて $x$ 以上になる添字を求めるなどすると問題を解けます．このとき，$x$ 以上にならない場合をケアする必要がありますが，findIndex はそういう設計になっていて Maybe Int が返ってくるので対応できます．最初に取り除いた $n - k$ 個の分を添え字に足す必要がありますが，Maybe Int の中身への加算なので fmap で足します．最後に fromMaybe -1 で Just なら結果に，Nothing なら -1 に変換できるので出力が得られます．
main = do [ n, k, x ] <- readInts as <- drop ( n - k ) . reverse . sort <$> readInts print $ fromMaybe -1 $ fmap ( ( n - k ) + ) $ findIndex ( x <= ) $ scanl (+) 0 as
別の実装として，降順ソートした $A$ の先頭 $n - k$ 要素を $0$ に置き換えてから累積和をとると他の部分が簡単になります．この変換は，$n - k$ 個の $0$ に無限個の $1$ が続くリストと zipWith (*) すれば実装できます
main = do [ n, k, x ] <- readInts as <- reverse . sort <$> readInts print $ fromMaybe -1 $ findIndex ( x <= ) $ scanl (+) 0 $ zipWith (*) as $ replicate ( n - k ) 0 ++ repeat 1
D
問題文 : https://atcoder.jp/contests/abc441/tasks/abc441_d
次数と $l$ の制約からパスの総数が $4^{ 10 }$ 以下であることが分かるので，DFS で潜って全部列挙してから適当に集計すればよいです．
ということで実装です．まずグラフのもち方ですが，頂点番号から辺リストへのアクセスが高速なら十分です．辺は行き先と重みのペアで表せるので，グラフは添字が Int で要素が [ ( Int, Int ) ] であるようなイミュータブル配列で表現できます．よって，入力で辺を表す行から読んだ $u, v, c$ から ( u, ( v, c ) ) という形にして accumArray ( flip (:) ) [] で畳み込めばいい感じになります．
DFS パートと集計パートは相互依存的になる気がしますが，集計パートを先に考えます．可能なすべてのパスに関して終点の頂点番号と重みの和からなるペアを並べたリストが得られると仮定します．このとき，これまた accumArray ( flip (:) ) [] で畳み込むと，頂点を添え字としてそれを終点とするパスの重みのリストが要素であるような配列を作れます．ここから assocs で添字と要素のペアからなるリストにして適当に filter して map fst すれば答えになります．filter する関数は snd をとってから $s$ 以上 $t$ 以下の判定をすればよくて，そういう要素が少なくともひとつあればよいということで any と組み合わせて（あまりラムダを使いたくない気持ちはありつつも）( any ( \w -> s <= w && w <= t ) . snd ) と実装できます．
DFS 部分は，
- グラフ
- 残り歩数
- 現在の頂点
- 現在までのパスの重み
を受け取って，その状態から作れるすべてのパスの情報を表すリストを返すようにします．グラフを g，現在地を u とすれば，g ! u についてリストの do を回すなどして適当に再帰して concat すれば実装できます．
main = do [ n, m, l, s, t ] <- readInts es <- replicateM m readInts let g = accumArray ( flip (:) ) [] ( 1, n ) do [ u, v, w ] <- es return ( u, ( v, w ) ) printList $ map fst $ filter ( any ( \w -> s <= w && w <= t ) . snd ) $ assocs $ accumArray ( flip (:) ) [] ( 1, n ) $ dfs g l 1 0 dfs _ 0 u d = [ ( u, d ) ] dfs g l u d = concat do ( v, w ) <- g ! u return $ dfs g ( pred l ) v ( d + w )
E
問題文 : https://atcoder.jp/contests/abc441/tasks/abc441_e
入力文字列を
\begin{align*}
\text{‘$A$'} &\mapsto 1 \\
\text{‘$B$'} &\mapsto -1 \\
\text{‘$C$'} &\mapsto 0
\end{align*}
という関数で map して得られる列を考えます．この列の累積和に着目すると，$1$ 以上であるような要素の個数が，先頭を左端点とする部分列の内で条件を満たすものの個数です．そういう要素の個数は，ヒストグラムを作っておくことでアクセスしやすくなります．また，元々の先頭以外を先頭とする部分列についても，元々の列から戦闘要素を取り除いたときに起こる変化をヒストグラムに反映することで同じ仕組みで計算できます．ということで，ヒストグラムは更新と部分和を効率的に計算できる形でもつ必要があって，競プロ的に代表的なデータ構造は Fenwick Tree です．Haskell 的には ac-library-hs の AtCoder.FenwickTree に実装があるので，これを使わせてもらうと楽です．
import qualified AtCoder.FenwickTree as Fenwick import qualified Data.Vector.Unboxed as VU -- テンプレ略 main = do n <- readInt s <- readStr let f 'A' = 1 f 'B' = -1 f 'C' = 0 ds = scanl1 (+) $ map f s as = elems $ accumArray (+) 0 ( -n, n ) $ map ( , 1 ) ds print $ sum $ runST do fenwick <- Fenwick.build $ VU.fromList as :: ST s ( Fenwick.FenwickTree s Int ) res0 <- Fenwick.sum fenwick ( n + 1 ) ( 2 * n + 1 ) ( res0 : ) <$> ( forM ( map ( + n ) ds ) $ \d -> do Fenwick.add fenwick d -1 Fenwick.sum fenwick ( d + 1 ) ( 2 * n + 1 ) )
F
問題文 : https://atcoder.jp/contests/abc441/tasks/abc441_f
つ ら い．
解法自体は 公式解説に任せますが，やるべきことは順方向と逆方向の DP をそれぞれやってから集計すると総括できます．ただ，Memory Limitt がかなりきついです．具体的に何が起こるかというと，STUArray で DP すると MLE して Mutable Vector で DP すると通るみたいなことが起きます（なんで？）．よく見ると 手癖で実装した C++ 解も 920 MB 使っていて結構あぶないです．もうちょっとゆるくても本質は損なわれなかったのでは……？
個人的には，STUArray の方をメインで使っているので（そして，そのまま書き換えたせいで一次元的なものを二次元的に使うことになったので）なかなか大変でした……．
import qualified Data.Vector as V import qualified Data.Vector.Unboxed as UV import qualified Data.Vector.Unboxed.Mutable as MV -- テンプレ略 main = do [ n, m ] <- readInts [ ps, vs ] <- transpose <$> replicateM n readInts let idx ( i, j ) = i * ( m + 1 ) + j solve ps vs = runST do dp <- MV.replicate ( ( n + 1 ) * ( m + 1 ) ) ( minBound `div` 2 ) :: ST s ( MV.MVector s Int ) MV.unsafeWrite dp ( idx ( 0, 0 ) ) 0 forM_ ( zip3 [ 0 .. ] ps vs ) $ \( i, p, v ) -> do forM_ [ 0 .. m ] $ \j -> do flip ( MV.unsafeModify dp ) ( idx ( i + 1, j ) ) =<< max <$> MV.unsafeRead dp ( idx ( i, j ) ) when ( j + p <= m ) do flip ( MV.unsafeModify dp ) ( idx ( i + 1, j + p ) ) =<< max . ( v + ) <$> MV.unsafeRead dp ( idx ( i, j ) ) forM_ [ 0 .. n ] $ \i -> do forM_ [ 1 .. m ] $ \j -> do flip ( MV.unsafeModify dp ) ( idx ( i, j ) ) =<< max <$> MV.unsafeRead dp ( idx ( i, j - 1 ) ) -- UV.unsafeFreeze dp dp1 = solve ps vs dp2 = solve ( reverse ps ) ( reverse vs ) ma = maximum do j <- [ 0 .. m ] return $ dp1 UV.! idx ( n, j ) f True False = 'A' f True True = 'B' f _ _ = 'C' putStrLn $ map ( uncurry f ) do ( i, p, v ) <- zip3 [ 0 .. ] ps vs let inc = or do j <- [ 0 .. m - p ] return $ dp1 UV.! idx ( i, j ) + dp2 UV.! idx ( n - 1 - i, m - p - j ) + v == ma exc = or do j <- [ 0 .. m ] return $ dp1 UV.! idx ( i, j ) + dp2 UV.! idx ( n - 1 - i, m - j ) == ma return ( inc, exc )
