ABC284 A～E+F をHaskellで

「Haskellで競技プログラミングをすると最初から容赦なくモナドが要求される」という聞き捨てならない呟きがあったので、getContentsとputStrLn だけで書く。

（というネタなので、元発言者の方はどうか気にされませんよう。）

1/11 F問題が解けたので追記。

A - Sequence of Strings

問題 ABC284A

reverse すればいい。

結果

main = getContents >>= putStrLn . unlines . abc284a . lines

abc284a :: [String] -> [String]
abc284a = reverse . tail

B - Multi Test Cases

問題 ABC284B

filter odd で残る数値の個数を数えればいい。

結果

chunksOf は元々 Data.List.Split の関数。

main = getContents >>= putStrLn . unlines . abc284b . lines

abc284b :: [String] -> [String]
abc284b = map (show . inner) . chunksOf 2 . tail

inner :: [String] -> Int
inner [ln,las] = length . filter odd . map read . words $ las

chunksOf n [] = []
chunksOf n xs = as : chunksOf n bs
  where
    (as,bs) = splitAt n xs

数を真面目にデコードしなくても、末尾の数字だけ見て判断することもできそう。

inner :: [String] -> Int
inner [ln,las] = length . filter (flip elem odds . last) . words $ las

odds = "13579"

import Data.Char

inner :: [String] -> Int
inner [ln,las] = length . filter (odd . digitToInt . last) . words $ las

C - Count Connected Components

問題 ABC284C

$N \leq 100$ なので少し大げさな道具に見えるが、UnionFind で連結成分を把握できる。
「正しい道具の使い方」としては、UnionFindし終わった後に、その分割の個数を数えればよいのだが、そこが割と面倒なので、unite操作が有効に行われた、つまり実際に連結が行われた回数を数えておいて、それを $N$ から引くことで、連結されずに残った分割の個数を得る。

結果

import qualified Data.IntMap as IM
import Data.List

main = getContents >>= putStrLn . unlines . abc284c . lines

abc284c :: [String] -> [String]
abc284c ls = map show [n - sum rs]
  where
    (n:m:_):uvs = map (map read . words) ls
    (_uf, rs) = mapAccumL step newUF uvs
    step uf (u:v:_) =
      case uniteUF uf u v of
        Nothing  -> (uf , 0)
        Just uf1 -> (uf1, 1)

-- @gotoki_no_joe
type UnionFind = IM.IntMap Int

newUF :: UnionFind
newUF = IM.empty

getRoot :: UnionFind -> Int -> (Int, Int)
getRoot uf i =
  case IM.lookup i uf of
    Nothing -> (i, 1)
    Just k | k < 0 -> (i, - k)
           | otherwise -> getRoot uf k

uniteUF :: UnionFind -> Int -> Int -> Maybe UnionFind
uniteUF uf i j
  | a == b    = Nothing
  | r >= s    = Just $ IM.insert a (negate $ r + s) $ IM.insert b a uf
  | otherwise = Just $ IM.insert b (negate $ r + s) $ IM.insert a b uf
  where
    (a, r) = getRoot uf i
    (b, s) = getRoot uf j

D - Happy New Year 2023

問題 ABC284D

効率的に素因数分解をする。

inner :: Int -> [Int]
inner n
  | a == b    = [a,c]
  | otherwise = [c,a]
  where
    (a:b:c:_) = primeFactors n

では、どうも時間が足らない。試し割り法では、線形探索が大きい方の素数まで行われることが原因だろうか。

素数の二乗で割り切れてその剰余が素数のとき、前者が $p$ で後者が $q$ という方法ではどうだろう。

inner :: Int -> [Int]
inner n = head $ catMaybes $ map f primes
  where
    f p =
      case divMod n (p * p) of
        (q,0) | isPrime q -> Just [p,q]
        _                 -> Nothing

あまり変わらない。

そうでなくて、$N$ は問題の制約を満たす数なので、素因数 $f$ が一つ見つかった時点で、それは $p$ か $q$ のいずれかである。なので $N / f$ が平方数 $k^2$ ならば $p = k, q = f$ さもなくば $p = f, q = N/p^2$ とわかる…だとまだ平方数判定とその根を求める計算が必要。$N$ が $f^2$ で割り切れるならば $p = f, q = N/p^2$ さもなくば $p = \sqrt{N/q}, q = f$ の方が楽か。

結果

整数平方根を求める計算は妥協した。

(1/9追記) 「妥協した」点に関する解説が公開されたのでリンクします。
⇒ 浮動小数点数オタクがAtCoder Beginner Contest 284のD問題をガチで解説してみる

main = getContents >>= putStrLn . unlines . abc284d . lines

abc284d :: [String] -> [String]
abc284d = map (unwords . map show . inner . read) . tail

inner :: Int -> [Int]
inner n =
  case divMod n (f * f) of
    (q, 0) -> [f, q]
    _      -> [round $ sqrt $ fromIntegral (div n f), f]
  where
    (f:_) = primeFactors n

-- @gotoki_no_joe
primeFactors :: Int -> [Int]
primeFactors n = loop n primes
  where
    primes = 2 : 3 : [y | x <- [5,11..], y <- [x, x + 2]]
    loop n pps@(p:ps)
      | n == 1    = []
      | n < p * p = [n]
      | r == 0    = p : loop q pps
      | otherwise = loop n ps
      where
        (q,r) = divMod n p

罠に見事に引っかかった感。

E - Count Simple Paths

問題 ABC284E

問題Cと入力の構造が同じだが、今回は問題のサイズが大きいので、入力は String でなく ByteString で処理することにする。

入力はこちらは普通にグラフとして扱うので、いつものように辺をノードごとに集計する。

パスの本数を数えるのに、長さ $L$ のパスがあったとき、途中下車するパスも合わせて全部で $L(L+1)/2$ 本と数えることになる。ノードを全て通る最長のパスがあったとき、それ一本だけで数える上限の $10^6$ を軽々と超えてしまう。

なので、打ち切りに気をつけつつ、また、$1 \to 2 \to 3 \to 4$ と $1 \to 3 \to 2 \to 4$ のような合流も意識して、「成長中のパスの末尾の頂点」と「パスに含まれるノードの集合」の対をキーとし、合流しているパスの本数を値とするマップを更新することで、幅優先探索をする。

abc284e :: [BS.ByteString] -> [String]
abc284e ls = map show [loop g 0 (M.singleton (1, IS.singleton 1) 1)]
  where
    (n:m:_):uvs = map (unfoldr (BS.readInt . BS.dropWhile isSpace)) ls
    g = accumArray (flip (:)) [] (1,n) [p | (u:v:_) <- uvs, p <- [(u,v),(v,u)]]
 
ub = 1000000
 
loop :: Array Int [Int] -> Int -> M.Map (Int, IS.IntSet) Int -> Int
loop g acc m
  | acc >= ub = ub
  | M.null m  = acc
  | otherwise = loop g acc1 m1
  where
    acc1 = acc + sum (M.elems m)
    m1 = M.fromListWith (+)
      [ ((i, IS.insert i s), c)
      | ((k, s), c) <- M.assocs m
      , i <- g ! k, IS.notMember i s
      ]

TLEした。
重複する計算を一度で済ませるために気を遣ったつもりが、逆に、Mapのキーの同一性を調べるのに時間がかかっているのだろうか。

結果

合流が起きてそのサイズが問題になるような入力では、すぐに上限に達して終わるだろうから、普通に深さ優先探索をすれば十分なのでは？ということで書いたら普通に間に合った。

深さ優先探索は、再帰呼び出しによるスタイルでなく、探索予定の状態を独自のスタックで管理する、ループによるスタイル（といってもHaskellではループも単なる末尾再帰で書くのだが）で書く。
というのは、再帰呼び出しによるスタイルだと、これまで数えたパスの本数を、再帰呼び出しの間をまたいで引き渡したり、上限に達したときに抜けたり、といった計算が表現しにくいから。

import qualified Data.ByteString.Char8 as BS
import Data.List
import Data.Char

import Data.Array
import qualified Data.IntSet as IS
import qualified Data.Map as M

main = BS.getContents >>= putStrLn . unlines . abc284e . BS.lines

abc284e :: [BS.ByteString] -> [String]
abc284e ls = [show ans]
  where
    (n:m:_):uvs = map (unfoldr (BS.readInt . BS.dropWhile isSpace)) ls
    g = accumArray (flip (:)) [] (1,n) [p | (u:v:_) <- uvs, p <- [(u,v),(v,u)]]
    ans = dfs g 0 [(1, IS.singleton 1)]

ub = 1000000

dfs :: Array Int [Int]    -- グラフ
    -> Int                -- パス数小計
    -> [(Int, IS.IntSet)] -- 探索スタック (末端ノード, 訪問済みノード集合)
    -> Int                -- 答え
dfs g acc _ | acc >= ub = ub
dfs g acc [] = acc
dfs g acc ((v,s):ts) = dfs g (succ acc) ts1
  where
    ts1 = foldr step ts $ filter (flip IS.notMember s) (g ! v)
    step v1 ts = (v1, IS.insert v1 s) : ts

入力を、A～D同様に String で処理した版もやってみた。

入力	実行時間	メモリ
ByteString	893ms	164MB
String	1889ms	141MB

本体の実行時間以上の1秒というオーバーヘッドが生じている。（が、まだ間に合った。）

F - ABCBAC

問題 ABC284F

解けたので追記。

結構 $T$ が長いので、入力を ByteString で扱う。
出力は、該当する $i$ が無い場合 -1 と一言、一方存在する場合は、$T$の前方 $i$ 文字と、後方 $N-i$ 文字を続けたもの。ByteString としてメモリ中にあるこれをわざわざ String にして出力するのももったいないので、出力側も ByteSting で扱う。

import qualified Data.ByteString.Char8 as BS

main = BS.getContents >>= BS.putStrLn . BS.unlines . abc284f . BS.lines

ByteStringの連結をするのももったいないので、とうとう、モナドな計算を利用する。
とはいえ mapM_ は Prelude 関数である。

import qualified Data.ByteString.Char8 as BS

main = BS.getContents >>= mapM_ BS.putStr . abc284f . BS.lines

$T$の$m$文字めから$n-1$文字め（0始まり）の区間を $T[m,n)$ と表すことにする。
逆転を $rev$ とする。

$S$は$T$に2度現れる。
まず順方向に、$i$ 文字で分断された前半$i$文字が$T$の先頭に、残りが末尾にある。
中央部に反転されて挿入される $S$ は、$N-1$ 文字目と$N$ 文字目の間を必ずまたいでいる。
つまり $T[0,i) = rev(T[N,N+i))$, $T[N+i,2N) = rev(T[i,N))$

素朴な方法

これを前からのんびり1文字ずつ比較すると、最後の方に違いが仕込まれた「いじわるデータ」に討ち取られるが、$T$ の $0,N-1,N,2N$ 文字めという4か所について先に調べる、という防衛をするだけで、TLEが2件だけという状態に持って行けたりする。が、結局間に合わない。

累積和？

列は変化しないので、a から z までの文字を別々に先頭から数えて何文字あるかをそれぞれの位置について数えて累積しておくと、任意の部分文字列 $T[a,b)$ について、そこに含まれる文字種とその個数が得られる。これが異なる場合は却下できる。

とやってみたが、数える仕事が重すぎて遅くなった上に、成分はわかってもその位置は不明なため、改善されなかった。

その変種として、区間に含まれるそれぞれの文字が奇数個かどうか、という1ビットだけを管理しても同様だった。

列の様子を反映するハッシュ

ラビン-カープ文字列検索アルゴリズムはハッシュを用いて文字列を検索する。
検索パターンに対するハッシュ値について、検索対称文字列の部分列に対してもハッシュ値を計算し、これが異なる場合は確実に文字列が異なることが判る、文字列が異なってもハッシュ値が衝突している場合もあるので、ハッシュ値が等しい場合だけ文字列を実際に比較する、というアプローチ。
ここで、検索対象文字列の先頭、検索パターンと同じ長さの部分列に対するハッシュ値を求めた後、1文字後ろにずれた区間のハッシュ値をゼロから求めるのではなく（そんなことをするなら文字で比較した方が早い）前のハッシュ値から高速に求められるような構造のハッシュ値を使う。

その一つのやり方が、文字種$K$文字数$N$の文字列を、$N$桁の$K$進数と見なす。そのままだと桁数が Int を容易に超えるので、大きな素数の剰余で考える、合同算術を用いる。
こうすると、ある文字列のハッシュ値が $h$ であるとき、その区間をひとつ右にずらして、区間の左端から追い出される文字が $c$ 、右から区間に入っていく文字が $d$ のとき、ずらした区間のハッシュ値は $hK - c \cdot K^N + d$ と求められる。

この問題に適用するには、中央部の逆転した $S$ と、両端に分断された $S$ の両方についてハッシュを求めることになる。
中央部は、$i=0$ のとき $T[0,N)$ なので、前半部に対してそのまま求める。動きとしては1文字ずつ後ろにずれるので、上の漸化式で計算する。
両端は、$i=0$ のとき $T[N,2N)$ で、こちらを逆順に求めておく。
こちらは、$i$ が増加するとき、$i$ 文字めを $c = T[i+N]$ から $d = T[i]$ に差し替えることになり、ハッシュ値は $h + (d - c) \cdot K^i$ で求められる。

結果

import qualified Data.ByteString.Char8 as BS
import Data.List
import Data.Char
import Data.Bits

main = BS.getContents >>= mapM_ BS.putStr . abc284f . BS.lines

abc284f :: [BS.ByteString] -> [BS.ByteString]
abc284f [nl,t]
  | null ans  = [BS.pack (show (-1)), cr]
  | otherwise = [BS.take i t, BS.drop (n+i) t, cr, BS.pack $ show i, cr]
  where
    Just (n,_) = BS.readInt nl
    bp = powerish 1 26 (pred n) -- 26^(n-1)
    ti i = ord (BS.index t i) - ord 'a'
    fpstep fp i = reg $ fp * 26 + ti i
    fL0 = foldl fpstep 0 [0..pred n]                         -- 前からN文字のhash
    fR0 = foldl fpstep 0 [n + pred n, pred n + pred n .. n]  -- 後からN文字のhash
    fLs = scanl flstep fL0 [0..pred n]
    flstep fp i = reg $ reg (fp - ti i * bp) * 26 + ti (n + i)
    fRs = scanl frstep fR0 $ zip [0..pred n] $ iterate (mul 26) 1
    frstep fp (i, bi) = reg $ fp + reg (ti i - ti (n + i)) * bi
    ans =
      [ i
      | (i, fl, fr) <- zip3 [0 .. pred n] fLs fRs
      , fl == fr
      , and [BS.index t j == BS.index t (    pred n + i - j) | j <- [0 .. pred i]]
      , and [BS.index t j == BS.index t (n + pred n + i - j) | j <- [i .. pred n]]
      ]
    i = head ans

cr = BS.singleton '\n'

modBase = 1000000007
reg x = mod x modBase
mul x y = reg (x * y)

-- @gotoki_no_joe
powerish i a b = foldl' {-'-} mul i [p | (True, p) <- zip bs ps]
  where
    bs = map odd $ takeWhile (0 <) $ iterate (flip div 2) b
    ps = iterate (\x -> mul x x) a