はじめに
実装メイン.テンプレは別記事.
A
問題文 : https://atcoder.jp/contests/abc439/tasks/abc439_a
与えられた整数 $n$ に対して $2^n - 2n$ を出力する問題です.これは言われた通りにやるだけで,
main = do n <- readInt print $ 2^n - 2 * n
のようにすれば AC できます.
これだけだと味気無いのでいつものようにワンライナー風味にしてみたいと思います.与えられた $n$ に対して (^2) と ( 2 * ) をそれぞれ適用してから引けばよいので,恒例の Control.Arrow にある (&&&) を使って $( 2^n, 2n )$ というタプルを作って,uncurry (-) に渡せば答えが求まります.
main = readInt >>= print . uncurry (-) . ( ( 2^ ) &&& ( 2 * ) )
B
問題文 : https://atcoder.jp/contests/abc439/tasks/abc439_b
各数字毎に考えると発散はしなさそうに思えます.なので,シミュレーションしてループしなければその内 $1$ になって Yes で,そうでなければ No です*1.この前提を認めるならば,あとは実装だけです.
まず,桁毎の二乗和を取る関数は標準っぽいところにある関数たちを使って
sum . map ( (^2) . digitToInt ) . show
という風に実装できます.あとは,ループ検出のために Data.IntSet などを使って既出を管理しつつ再帰することで問題を解くことができます.
main = readInt >>= putStrLn . yesno . solve ISet.empty solve _ 1 = True solve s n | ISet.member n s = False | otherwise = solve ( ISet.insert n s ) $ f n where f = sum . map ( (^2) . digitToInt ) . show
C
問題文 : https://atcoder.jp/contests/abc439/tasks/abc439_c
$n$ 以下の平方数の個数は $O( \sqrt n )$ 個であり,$x, y$ の候補も $O( \sqrt n )$ 個ずつなので組合せは $O( \sqrt n \times \sqrt n ) = O( n )$ 通りです.従って,余計な範囲を探索しないように二重ループ(に相当することを)すれば問題を解けます.具体的には,$x^2 \leq n$ であるような各 $x$ に対して,$x^2 + y^2 \leq n$ であるような $y$ を全て試し,得られる整数それぞれについて個数をカウントします.
集計部分の方が面倒くさいという説がわたしの中にはありますが,上記の条件を満たす $x, y$ によって生成される $x^2 + y^2$ からなるリストに対して map ( , 1 ) して得られるリストを使って accumArray (+) 0 することで集計できます.あとは,assocs で添字と要素のタプル列にしてから適当に加工することで,出力するべきリストを構成できます.
main = do n <- readInt let res = map fst $ filter ( ( == 1 ) . snd ) $ assocs $ accumArray (+) 0 ( 0, n ) $ concatMap ( map ( , 1 ) ) do x <- takeWhile ( ( <= n ) . (^2) ) [ 1 .. ] let xx = x^2 return do yy <- takeWhile ( ( <= n ) . ( + xx ) ) $ map (^2) [ x + 1 .. ] return $ xx + yy print $ length res printList res
D
問題文 : https://atcoder.jp/contests/abc439/tasks/abc439_d
まず,$\max( i, j, k ) = j$ の場合について解ければ,入力 $A$ を反転した列に対して同じアルゴリズムを適用することで(元々の列における添字で)$\min( i, j, k ) = j$ の場合も解けます.なので $\max( i, j, k ) = j$ の場合について考えます.また,
\[
A_i : A_j : A_k = 7 : 5 : 3
\]
ということは,ある整数 $x$ が存在して,
\begin{align*}
A_i &= 7x \\
A_j &= 5x \\
A_k &= 3x
\end{align*}
ということです.今は $\max( i, j, k ) = j$ の場合について考えているので,列を先頭から走査すると考えると最も先行している添字は $j$ です.上述のような整数 $x$ の存在性も踏まえると,$A_j \bmod 5 = 0$ でなければその $j$ に対して valid な $i, k$ は有りません.そうでない場合,
\begin{align*}
p &= \#\left\{ i \in \{ 1, 2, \dots, j - 1 \} \mid A_i = \frac{ A_j }{ 5 } \times 7 \right\} \\
q &= \#\left\{ k \in \{ 1, 2, \dots, j - 1 \} \mid A_k = \frac{ A_j }{ 5 } \times 3 \right\}
\end{align*}
としたとき,valid な $i, k$ の組合せが $pq$ 個ということなります.ここで,集合 $S$ に対して $\#S := |S|$ です.
ということで,$A$ を先頭から走査しながら,$\frac{ A_j }{5} \times 7, \frac{ A_j }{5} \times 3$ の個数を得られればよいです.よって,各時点までの値の出現回数を Data.IntMap などで数えることで,$\max( i, j, k ) = j$ の問題を解くことができ,$A$ を反転した列に対して同じことをして結果を足し合わせることで元々の問題を解けます.
実装については最終的に「Hoogle で IntMap のインターフェイスを調べよう!」ということになる気がしますが,敢えて書くなら insertWith を使うとキーが存在しない場合の代入と,存在する場合の更新を一箇所で書けてよい気がします.
main = do n <- readInt as <- readInts print $ solve as + solve ( reverse as ) solve as = sum $ snd $ mapAccumL solve' IMap.empty as where solve' s a = ( IMap.insertWith (+) a 1 s, r ) where r = if a `mod` 5 == 0 then product $ map ( fromMaybe 0 . ( s IMap.!? ) ) [ a `div` 5 * 3, a `div` 5 * 7 ] else 0
E
問題文 : https://atcoder.jp/contests/abc439/tasks/abc439_e
詳しい説明は公式解説等に投げますが,$A, B$ を zip した列 $\langle ( A_1, B_1 ), ( A_2, B_2 ), \dots, ( A_n, B_n ) \rangle$ を
- $A_i < A_j$ ならば $i$ 番目の項が手前
- そうでないとき,$B_i < B_j$ ならば $j$ 番目の項が手前
という風にソートしてから $B$ のみを取り出した(map snd した)列の最長増加部分列 (Longest Increasing Subsequence; LIS) の長さが答えになります.
LIS を求める方法はいくつかあるかと思いますが,ここでは
\[
\mathit{ dp }[b] = \text{$b$ を末尾とする増加部分列の長さの最大値}
\]
という DP を考えます.この DP をそのまま実装しようとすると入力列の要素の値が大きすぎて困ってしまいますが,座標圧縮と呼ばれる変換を施すことで対応できます.本稿では,実装を簡略化するために $1$-based で座標圧縮します.DP 配列の初期化は
\[
\mathit{ dp }[b] = \begin{cases}
0 & \text{($b = 0$)} \\
-\infty & \text{(otherwise)}
\end{cases}
\]
となります.更新は,前述のソートおよび座標圧縮済みの列の要素を先頭から走査し,各要素 $b$ に対して
\[
\mathit{ dp }[b] \leftarrow \max\{ \mathit{ dp }[0], \mathit{ dp }[1], \dots, \mathit{ dp }[ b - 1 ] \} + 1
\]
となります.$\max$ を求めるときに毎回 $\Theta( n )$ 時間かけるともちろん TLE しますが,DP 配列自体を Segment Tree に載せて $O( \log n )$ 時間で処理できるようにすることで TL に間に合います.
必要な部品が揃ったので実装の話に入ります.$A$ の昇順・$B$ の降順になるように列というか添字をソートする部分は zip as bs を sortBy してから map snd することで実装できます.sortBy は Ordering を返すようなインターフェースになっていますが,(個人的には忘れがちなのですが)compare が $2$ 要素を比較して Ordering を返す関数です.ということで
\( a1, b1 ) ( a2, b2 ) -> if a1 == a2 then compare b2 b1 else compare a1 a2
という関数でソートすれば望みのものが得られます.または,$A$ で昇順・$B$ で降順ということを活用してタプル $( a, b )$ を $( a, -b )$ と見做してソートしても同じ結果になるので,
\( a, b ) -> ( a, -b )
という関数を使って sortOn してもよいです.記述量的にはこっちの方が短くて楽ですね.
次は座標圧縮ですが,典型的には元の列をソートしてから重複を取り除いた列に対する二分探索で変換後の値を決める実装がなされます.二分探索は ac-library-hs の AtCoder.Extra.Bisect に実装があるのでこれを使わせてもらいます.入力を Data.Vecotr で受け取ることになっているので,どうせ Data.Vector にするなら(隣接する要素の)重複除去は Data.Vector にある uniq を使えばよいです(リストだけでやる場合は sort して group して map head する実装をよくやります).ということで,ソート・重複除去済みの Data.Vector を使って upperBound をすることで $1$-based での座標圧縮が完了します.
DP 部分には Segment Tree が欲しいと書きましたが,これまた ac-library-hs に AtCoder.SegTree があるのでこれを使わせてもらいます.今回は $\max$ を載せたいわけですが,Data.Semigroup に Max a があるので Max Int を載せるだけで望みの Segment Tree が手に入ります.あとは,上述のような DP を実装するだけです.
import qualified Data.Vector as V import AtCoder.Extra.Bisect import Data.Semigroup import qualified AtCoder.SegTree as SegTree -- テンプレ略 main = do n <- readInt [ as, bs ] <- transpose <$> replicateM n readInts let uniqBs = V.uniq $ V.fromList $ sort bs -- bs' = map ( upperBound uniqBs . snd ) $ sortBy ( \( a1, b1 ) ( a2, b2 ) -> if a1 == a2 then compare b2 b1 else compare a1 a2 ) $ zip as bs bs' = map ( upperBound uniqBs . snd ) $ sortOn ( \( a, b ) -> ( a, -b ) ) $ zip as bs print $ getMax $ runST do st <- SegTree.new ( n + 1 ) :: ST s ( SegTree.SegTree s ( Max Int ) ) SegTree.write st 0 0 forM_ bs' $ \b -> do SegTree.write st b =<< succ <$> SegTree.prod st 0 b SegTree.allProd st
F
問題文 : https://atcoder.jp/contests/abc439/tasks/abc439_f
解法自体は別記事にしてあるのでそちら(か公式解説等)をご参照頂くとして,早速実装の話をします.
こちらの問題についても ac-library-hs が大活躍で,${} \bmod {998{,}244{,}353}$ も Fenwick Tree も ac-library-hs にあるのでこれらを使います.そうすると上述の別記事の内容を実装するだけという話になって書くことがあまり無いような気がしますが,強いて言えば,いくつかの部分を ($), <*> を使って Applicative スタイルにすると余計な変数束縛を消せてよいかなという気がします.例えば,$L_i \times R_i$ を求めるところは $L, R$ をミュータブル配列にした場合は参照がアクションになってしまいますが,
(*) <$> readArray left i <*> readArray right i
という風に書けたりします.
import AtCoder.ModInt import Data.Monoid import qualified AtCoder.FenwickTree as FT type Mint = ModInt998244353 -- テンプレ略 main = do n <- readInt ps <- listArray ( 0, n - 1 ) . map pred <$> readInts print $ runST do fenwickL <- FT.new n :: ST s ( FT.FenwickTree s ( Sum Mint ) ) fenwickR <- FT.new n :: ST s ( FT.FenwickTree s ( Sum Mint ) ) left <- newArray ( 0, n - 1 ) 0 :: ST s ( STArray s Int Mint ) right <- newArray ( 0, n - 1 ) 0 :: ST s ( STArray s Int Mint ) forM_ [ 0 .. n - 1 ] $ \i -> do FT.add fenwickL ( ps ! i ) 1 writeArray left i =<< getSum <$> FT.sum fenwickL 0 ( ps ! i ) forM_ [ n - 1, n - 2 .. 0 ] $ \i -> do FT.add fenwickR ( ps ! i ) 1 writeArray right i =<< getSum <$> FT.sum fenwickR 0 ( ps ! i ) res1 <- forM [ 0 .. n - 1 ] $ \i -> do (*) <$> readArray left i <*> readArray right i let f t ( l, r ) = ( t * 2 + l, t * r ) res2 <- snd . mapAccumL f 0 <$> ( zip <$> getElems left <*> getElems right )
*1:よりちゃんとした証明は公式解説等を参照のこと