ABC362をHaskellで

A - Buy a Pen

問題 ABC362A

シグネチャを決める。

abc362a :: [Int]  -- R,G,B
        -> String -- C
        -> Int    -- 答え

先頭1文字で捨てる選択肢を判断し、残りの2つのうち安い方を選ぶ。

結果

abc362a [r,g,b] (c:_) =
  case c of
    'R' -> min g b
    'G' -> min r b
    'B' -> min r g

B - Right Triangle

問題 ABC362B

シグネチャを決める。横着する。

abc362b :: [[Int]]  -- Xi,Yi
        -> Bool     -- 答え

• 2辺のベクトルを選んで内積をとって0になることで直交を調べる
• 辺の長さが三平方の定理を満たすか調べる

のどちらでもいいだろう。

内積

abc362b [p1,p2,p3] = any (0 ==) $ zipWith dp [v1,v2,v3] [v2,v3,v1]
  where
    [v1,v2,v3] = zipWith (zipWith (-)) [p1,p2,p3] [p2,p3,p1]
    dp v1 v2 = sum $ zipWith (*) v1 v2

コードアシストが「any (0 ==) はダセーから elem 0 にしろ」と。
同じだけど違うんだなぁ。

三平方の定理

$a^2 + b^2 = c^2$ は $a^2 + b^2 + c^2 = 2c^2$

abc362b [p1,p2,p3] = sum ds == 2 * maximum ds
  where
    dist2 p1 p2 = sum $ map (^2) $ zipWith (-) p1 p2
    ds = zipWith dist2 [p1,p2,p3] [p2,p3,p1]

C - Sum = 0

問題 ABC362C

シグネチャを決める。
答えがないときは空リストで表すことにする。

abc362c :: Int      -- N
        -> [[Int]]  -- Li,Ri
        -> [Int]    -- 答え

$L_i$ の総和が0より大きいとき、全ての項で可能な限り小さい値を選んでも結果は0より大きくなってしまう。$L = \sum L_i$ とする。
$R_i$ の総和が0より小さいとき、全ての項で可能な限り大きな値を選んでも結果は0より小さくなってしまう。
そうでないとき、うまく調節すれば総和を0にできる。

「うまく」とは：
仮に、全ての項で最も小さい値を選ぶ。総和は $L \leq 0$ となり、$-L$ 足らない。
この不足分は、補えるところで任意に補えばいいので、順に、
第 $i$ 項で補える最大値 $R_i - L_i$ と、現在の残りとの小さい方を補填し、
補填した分で残りの現在値を更新する、を繰り返せばよい。

結果

残り $rest$ をいきなり全額補填すると $L_i + rest$ になるが、$R_i$ が上限なので、その小さい方を出力する、という流れで計算している。

abc362c :: Int -> [[Int]] -> [Int]
abc362c n lrs
  | 0 < accL || accR < 0 = []
  | otherwise = snd $ mapAccumL step (negate accL) lrs
  where
    accL = sum $ map head   lrs
    accR = sum $ map (!! 1) lrs
    step rest (l:r:_) = (rest1, out)
      where
        out = min r (l + rest)
        rest1 = rest - (out - l)

D - Shortest Path 3

問題 ABC362D

シグネチャを決める。

abc362d :: Int      -- N
        -> Int      -- M
        -> [Int]    -- Ai
        -> [[Int]]  -- Ui,Vi,Bi
        -> [Int]    -- 答え

辺の重みだけならダイクストラ法で求められる。
頂点の重みを辺の重みに置き換えるために、元のグラフの頂点 $i$ を、そこに入る有向辺が届く $i_{IN}$ と、そこから出て行く辺が出発する $i_{OUT}$ に分割し、元の無向辺はそのように張る。
さらに、$i_{IN}$から$i_{OUT}$へ重さ$A_i$な辺を張る。
問題の要求は、この新しい有向グラフにおける、$1_{IN}$ から各 $i_{OUT}$ への距離と読み替えられる。

結果

$i_{IN}$を$i$、$i_{OUT}$を$i+N$ に割り振った。

import Control.Monad.ST
import Data.Array.ST
import Data.Array

abc362d :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc362d n _m as uvbs = runST $
  do
    dists <- dijkstra (n + n) (edge !) 1
    forM [n + 2 .. n + n] (readArray dists)
  where
    edge = accumArray (flip (:)) [] (1, n + n) $
           [(i, (i + n, a)) | (i, a) <- zip [1..] as] ++
           [p | u:v:b:_ <- uvbs, p <- [(n + u, (v, b)), (n + v, (u, b))]]

-- ダイクストラ法
dijkstra :: Int                  -- 頂点数N (頂点1～N)
         -> (Int -> [(Int,Int)]) -- 隣接頂点とその辺の重み、グラフの情報
         -> Int                  -- 開始点
         -> ST s (STUArray s Int Int)
-- 実装略

提出はこちら

公式解説のやり方

上の、頂点を分割する方法はフレンズさんのヒントと同じ考え方。

公式解説では、頂点を増やさずに、全ての入り辺の重みに自分の頂点の重みを足し込んだ有向グラフを作る、結果には$A_1$の重みをさらに足す、という解法をとっていた。
これなら頂点が増えないので、計算はもっと速いはずだ。

import Control.Monad.ST
import Data.Array.ST
import Data.Array

abc362d :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc362d n _m as uvbs = runST $
  do
    dists <- dijkstra (n + n) (edge !) 1
    forM [2 .. n] (\i -> (head as +) <$> readArray dists i)
  where
    a = listArray (1,n) as
    edge = accumArray (flip (:)) [] (1, n) $
           [ p
           | u:v:b:_ <- uvbs
           , p <- [ (u, (v, b + a ! v))
                  , (v, (u, b + a ! u))]]

	時間(ms)	メモリ(MB)
頂点2倍	1668	177
公式	1127	166

速かった。

E - Count Arithmetic Subsequences

問題 ABC362E

シグネチャを決める。

abc362e :: Int      -- N
        -> [Int]    -- Ai
        -> [Int]    -- 答え

長さ1の数列は、その項の値によらず、等差数列といえる。
長さ2の数列も、その2項の値によらず、等差数列といえる。

2項あって初めて公差が定まるので、数列 $A$ からまず任意の2項を選択し、その公差を維持して数列を延長する方法の数を数えることで、長さ3以上の等差数列の個数を求めることができるだろう。

数列の項を前から確認していき、$A_i$ を見たとき、

• それより手前の全ての項 $A_j$ ($1 \leq j < i$) と続けて、長さ2の等差数列ができる。続きの項で長さ3以上にできるとき数えられるように、諸情報を記録しておく。
• それより手前の項だけで作られる長さ $k$ ($2 \leq k$) 公差 $D$ 末尾 $L$ な数列について、$L+D=A_i$ ならば、長さ $k+1$ 公差 $D$ 末尾 $A_i$ な数列を作ることができるのでこれを数える。これより後の項 $A_m$ ($i < m, A_i = A_m$) でふたたび類似の数列が作れるので、上書きではなく追加になる。

数列の集合から $L+D=A_i$ を満たすものを探して取り出すのは難儀そうなので、

• 次に来て欲しい値 $L+D$ をキーとする IntMap に格納することで、lookup だけで取り出せるようにする。
• 公差や長さが一致するような数列も多々現れると思われるので、長さと公差の対をキーにし、その個数を内容とする Map を、上の IntMap の内容とする。

というデータ構造

type Dat = IntMap -- キーは、次に数列で出現してほしい値
             (Map ( Int  -- キーは、数列の長さと
                  , Int) -- 　　　　公差の組
                  Int   -- 内容は、そのような数列の個数
             )     -- というMapが内容

とする。
$A_i$ を見たとき、IntMap を検索して出てきた Map に関して、
その全ての対応 $((長さ k, 公差 D), 個数 c)$ に対して、
「次に来て欲しい値 $A_i + D$, 長さ $k+1$, 公差 $D$」な数列の個数に $c$ を足し込む。

結果

import Data.List
import qualified Data.IntMap as IM
import qualified Data.Map as M
import Data.Array

type State = (IM.IntMap Int, Dat)

abc362e :: Int -> [Int] -> [Int]
abc362e 1 _ = [1]
abc362e 2 _ = [2, 1]
abc362e n as = n : div (n * pred n) 2 : elems ans
  where
    (_,datN) = foldl' step (IM.empty, IM.empty) as
    step :: State -> Int -> State
    step (cnt1, cntk) ai = (cnt11, cntk1)
      where
        cnt11 = IM.insertWith (+) ai 1 cnt1
        len2 = [ (ai + d, [((2, d), cnt)])
               | (aj, cnt) <- IM.assocs cnt1, let d = ai - aj ] -- 新規の [Aj,Ai]
        m = IM.findWithDefault M.empty ai cntk
        len3 = [ (ai + d, [((succ len, d), cnt)])
               | ((len, d), cnt) <- M.assocs m ] -- 新規の [...,Ai]
        cntk1 = IM.unionWith (M.unionWith add) cntk $
                IM.map (M.fromListWith add) $ IM.fromListWith (++) $ len2 ++ len3
    ans = accumArray add 0 (3,n)
          [(k, c) | m <- IM.elems datN, ((k, _D), c) <- M.assocs m, k > 2]

add :: Int -> Int -> Int
add x y = mod (x + y) 998244353

$A_i$ を末尾とするものだけの新たな Dat を作って、これを直前の cntk と足し合わせる計算を
IM.unionWith (M.unionWith add) で書けるところがいいと思う。

フレンズさんのヒント

アライグマ いわく

「E問題はDPなのだ！
DP[i][k][d]=末尾がA[i]で、長さがkで、公差がdの等差数列の個数
でDPすればO(N^3)で解けるのだ！

公差は最大で $10^9 - 1$ になるので、そのまま配列の添え字にするのは厳しい。
出現する公差は $A_i - A_j$の値 $N(N-1)/2$ とおりしかないので、それを先に列挙して、背番号で呼ぶ。

j を振るとそれに応じて di が決まり、それを添え字として配る必要があるので、一段の遅延配列DPでなく、iとkだけを添え字とする外側の配列と、di を添え字とする内側の配列の二段構えにした。

import Data.List
import qualified Data.IntMap as IM
import Data.Array

abc362e :: Int -> [Int] -> [Int]
abc362e 1 _ = [1]
abc362e n as = n : elems ans
  where
    dimax = IM.size dindex
    dindex :: IM.IntMap Int
    dindex = IM.fromList [(d, dindexf d) | a:bs <- tails as, b <- bs, let d = b - a]
    dindexf d = maybe 1 (succ . snd) $ IM.lookupLT d dindex

    a = listArray (1,n) as

    bnds = ((1,2),(n,n))
    cnts = listArray bnds $ map cntsf $ range bnds
    cntsf (i,k) = accumArray add 0 (1, dimax)
      [ (di, if k == 2 then 1 else cnts ! (j, pred k) ! di)
      | j <- [1 .. pred i], let di = dindex IM.! (a ! i - a ! j) ]

    ans = accumArray add 0 (2, n)
          [(k, c) | ((_i, k), dic) <- assocs cnts, c <- elems dic]

構造はこっちの方がシンプルか。
$O(N^3)$な表を何度も舐めては長さを１ずつ伸ばしていくので、全体で$O(N^4)$になる気がするのだけど、どうだろう。

追記

フレンズさんの方法は配るDPで考えるべきもので、すると、
ある $i$ から $j$ ($i < j$) へ配るとき、交差 $A_j - A_i$ であるようなものだけ、長さ $k$ ($2 \leq k \leq N$) の値を配るので、計算量は$O(N^3)$であってる。

そのように命令型で書いてみた。
cnts IM.! d ! (k,i) が DP[i][k][d] に相当する。

abc362e :: Int -> [Int] -> [Int]
abc362e 1 _ = [1]
abc362e n as = (n :) $ runST $
  do
    arrs <- forM (IS.elems dset) (\_ -> newArray ((2,1), (n,n)) 0) :: ST s [STUArray s (Int,Int) Int]
    let cnts = IM.fromDistinctAscList $ zip (IS.elems dset) arrs
    forM_ ijds (\(i, j, d) -> do
      let cntd = cnts IM.! d
      modifyArray cntd (add 1) (2, j)
      forM_ [3 .. n] (\k -> do
        x <- readArray cntd (pred k, i)
        modifyArray cntd (add x) (k, j)
        )
      )

    ans <- newArray (2, n) 0 :: ST s (STUArray s Int Int)
    forM_ (IM.elems cnts) (\cntd -> do
        forM_ [2 .. n] (\k -> do
          acc <- sum <$> forM [2 .. n] (\i -> readArray cntd (k, i))
          modifyArray ans (add acc) k
          )
      )
    getElems ans
  where
    ijds = [(i,j,aj - ai) | (i,ai):jajs <- tails $ zip [1..] as, (j,aj) <- jajs]
    dset = IS.fromList [d | a:bs <- tails as, b <- bs, let d = b - a]
    modifyArray arr fun idx = readArray arr idx >>= writeArray arr idx . fun

add :: Int -> Int -> Int
add x y = mod (x + y) 998244353

F - Perfect Matching on a Tree

問題 ABC362F

シグネチャを決める。横着する。

import Data.Array

abc362f :: Int      -- N
        -> [[Int]]  -- ui, vi
        -> [[Int]]  -- 答え
abc362f n uvs = ...
  where

わからないので解説を見た。

仮の根を決めて部分木の大きさを数える $O(N)$
それを頼りに重心を探す $O(N)$
重心の部分木に着目して、ペアのノードは違う部分木から選べばどう選んでも最長になる
ただし、頂点が偶数のときは、部分木の頂点だけだと奇数個になって一つ余るので、それは重心とペアにする

公式解説では、最後のステップをおおむね言葉通りにやる話。
部分木の頂点数は$N/2$以下なので、部分木ごとの頂点リストを全て連結したものを二等分し、前から順に選べば、同じ部分木から選んでしまう事故は起きない。
これを折り返したりすると、前半と後半に跨いだひとつの部分木から両者を選ぶ事故が起きる。

ユーザ解説 by MMNMMはさらに上手い。
木を行きがけ順、または帰りがけ順にした頂点リストは、スタート地点を変えても、そのリストの回転にしかならない。（兄弟の選択順についてはうまいことやるとして）
なので、この頂点リストを二等分したリストは、重心の一つの部分木に属する頂点が跨いでいる可能性があるが、その頂点数は$N/2$以下なので、どちらも前から選べば、開始点がどこかに関わらず、事故ることはない。
頂点が奇数のときに、二等分する前に重心を除いておくことだけが必要になり、頂点が偶数のときは重心を探す必要すらない（！）

考える

グラフの配列をいつものように作る。

    g = accumArray (flip (:)) [] (1,n) $ concat [[(u,v),(v,u)] | u:v:_ <- uvs]

木の深さ優先走査

• 行きがけ順の頂点リストを作る
• 重心を探すために、仮頂点から見た部分木のサイズを数える

の2回、深さ優先の木の走査をするので、その両者を抽象化する関数を立ててみる。

depthFirstTraverse func ini graph root は、
頂点 i からの出辺リスト graph ! i で表される無向グラフとしての木に
root を根とした深さ優先走査を行う
頂点 v に降りてきたときの累積情報 acc0 と、
v の子を回って戻ってきた後の累積情報 acc1 について
func v acc1 acc2 した結果が v の親に戻される
root に降りてくるときの acc の初期値を ini で与える

depthFirstTraverse :: Ix i => (i -> a -> a -> a) -> a -> Array i [i] -> i -> a
depthFirstTraverse func ini graph root = loop root root ini
  where
    loop parent v acc =
      func v acc $ foldr (loop v) acc $
      [c | c <- graph ! v, c /= parent]

関数を差し込む位置やその引数は、今回の応用に必要なように入れただけなので、これが完璧な汎用性を持つ正解かはわからない。

全ての部分木の大きさを数える

1を仮の根とする。

    size1 = array (1,n) $ snd $ depthFirstTraverse sizeFun (0, []) g 1
    sizeFun v (t0, _) (t1, ps) = (t2, (v, t2 - t0) : ps)
      where
        t2 = succ t1

v に降りてきた時刻が t0、子を巡って戻ってきたときの時刻が t1 で、v を抜ける瞬間にさらに1時間が進む。この時刻の差が v の部分木の頂点数なので、(v, サイズ) 対にそれを追加して返す。

重心を探す

仮の根1から出発して、子のサイズが全て$N/2$以下になる頂点を探す。最大のものがそれを越えていたら、その子に降りていけばいい。

    n2 =div n 2
    centro = findCentroid 1 1
    findCentroid p v
      | smax <= n2 = v
      | otherwise = findCentroid v c
      where
        (smax, c) = maximum [(size1 ! c, c) | c <- g ! v, c /= p]

これは depthFirstTraverse を使えなかった。

行きがけ順の頂点リストを作る

1から始める。

    nodes = depthFirstTraverse nodesFunc [] g 1
    nodesFunc v _ vs = v : vs

子を巡って戻ってきたときの頂点リストより前に v がリストに載っているはずである。
帰りがけ順にしたいときは、nodesFunc v _ vs = vs ++ [v] はダサいので、depthFirstTravers の高階引数が足りていないということになる。

paranoicDepthFirstTraverse f1 f2 f3 ini graph root = loop root root ini
  where
    loop parent v acc0 =
      f3 v acc0 $
      foldr (\c acc -> f2 v acc0 c acc $ loop v c acc) (f1 v acc0) $
      [c | c <- graph ! v, c /= parent]

ここまで関数を差し込めるようにすれば

    nodes = paranoicDepthFirstTraverse nf1 nf2 nf3 [] g 1
    nf1 v acc0 = v : acc0
    nf2 _v _acc0 _c _acc = id
    nf3 _v _acc0 = id

とできるが、多分やり過ぎだろう。

半分にしてくっつける

abc362f n uvs = zipWith (\a b -> [a,b]) as bs
  where
    (as,bs) = splitAt n2 $ if odd n then filter (centro /=) nodes else nodes

$N$が偶数のとき、遅延評価でcentroの計算自体がスルーされる。

どこから巡っても同じなので、$N$が奇数のときはせっかくだから centro から出発し、先頭にあるそれを tail で消して終わりでもいいか。

    nodes
      | odd n = tail $ depthFirstTraverse nodesFunc [] g centro
      | otherwise =    depthFirstTraverse nodesFunc [] g 1
    (as,bs) = splitAt n2 nodes

結果：AC 394ms 111MB

G - Count Substring Query

問題 ABC362G

シグネチャを決める。
実際には入力が大きいので ByteString を使う。

abc362g :: String   -- S
        -> Int      -- Q
        -> [String] -- Ti
        -> [Int]    -- 答え
abc362g s q ts = ...

さっぱりわからないが、僕のTLにはこんな有用なtweetが流れてくる：

Aho-Corasick法

ググるといろいろ解説があるが読んでもよくわからないので、原典にあたる。
Alfred V. Aho and Margaret J. Corasick. 1975. Efficient string matching: an aid to bibliographic search. Commun. ACM 18, 6 (June 1975), 333–340.

アルゴリズム2は、Trie木に相当するgoto関数を作る。
状態番号0を初期状態とし、それぞれ、入力文字に応じて遷移先の状態番号を割り当てる。
遷移先がない場合には、特別な記号failを割り当てる。実装上は-1で表す。
なお状態0については、failとせず0に遷移させる。

また同時に、根から始めて $T_i$ が受理された頂点に印を付けるoutput関数の構築の前半を行う。
これは無理に同時にやらなくても構わない。

元論文ではgotoの配列が可変長になっているのが微妙にうざい。C++のvectorがうらやましい。
$1 + \sum |T_i|$ が上限なので、一旦それで配列を張っておくしかないか。
ks の先頭の文字で分類し、再帰的に状態を割り振る。

import Data.Array.Unboxed
import qualified Data.Sequence as Q

buildGoto :: [String] -> UArray (Int, Char) Int
buildGoto ks = goto
  where
    stmax = sum $ map length ks

    goto = array ((0,'a'), (stmax,'z')) $ recur 1 $ Q.singleton (0, ks)
-- 引数：次に割り当てる状態番号
-- キュー引数：状態番号、そこで分岐するキーワードリスト
-- 結果：arrayに与える ((状態番号, 文字), 遷移先状態番号) のリスト
    recur _ Q.Empty = []
    recur st ((i, ks) Q.:<| iks) =
        [ ((i, x), j) | (j, (x, _)) <- zip [st ..] succs] ++
        [ ((i, x), fail) | (x, []) <- assocs arr] ++
        recur (st + length succs) (iks Q.>< Q.fromList [(j, xss) | (j, (_, xss)) <- zip [st ..] succs])
      where
        fail = if i == 0 then 0 else -1
        arr = accumArray (flip (:)) [] ('a','z') [(x, xs) | x:xs <- ks] -- 先頭の文字で分類
        succs = filter (not . null . snd) $ assocs arr -- 子

残りの仕事を Data.Sequence を用いたキューに入れることで、ノード番号の順がノードの低い順に割り当てられる。
すると、以降の計算において、元論文ではキューを用いた木の探索をしている部分を、単純に添え字に関するforループに置き換えられる。
（ただしこの節の実装では遅延配列で対処しているのでループ自体が不要になっている。）

outputは、構築されたgotoに対して、キーワードで状態0から遷移した先を記録することで、アルゴリズム2の段階の内容が構築できる。
後で、outputの値をキーワードの背番号リストにしていたことがメモリ異常消費の原因であったことがわかったので、この段階で IntSet にしておく。
ついでに、これらの配列の添え字の最大値も求める。buildGotoのrecurはそれを知っているが、副作用のない計算ではそれを取り出すのは難しい。

import qualified Data.IntSet as IS

buildOutput :: UArray (Int, Char) Int -> [String] -> (Int, Array Int IS.IntSet)
buildOutput goto ks = (ub, out1)
  where
    run k = foldl' step 0 k
    step s c = goto ! (s, c)
    ub = maximum $ map run ks
    out1 = accumArray IS.union IS.empty (0, ub)
           [(run k, IS.singleton i) | (k,i) <- zip ks [1..]]

続くアルゴリズム3で、失敗関数 f を構築し、同時に output を加工している。これも同時にする必要はない。

キューだといいつつ集合の記号が使われている。本質は、状態0からの高さが低い方から計算する必要があることにある。これは、f関数が必ず木の根に近い方に飛ぶため、根の方から計算することで、確定済みの値を順次利用できるからである。
遅延配列を使えばそういう面倒なしに宣言的に書ける。

buildFail :: Int -> UArray (Int,Char) Int -> Array Int Int
buildFail ub goto = f
  where
    f = array (1, ub)
        [ (s, if r == 0 then 0 else st)
        | ((r, a), s) <- assocs goto, s > 0
        , let st = head $ dropWhile (0 >) $ map (\s -> goto ! (s, a)) $ tail $ iterate (f !) r ]

アルゴリズム3のendでない最後の行で output の拡張をしている。
f 関数の全ての対応 $f(s) = fs$ について、$output(s)$ に $output(fs)$ を追加する。

extendOutput :: Array Int IS.IntSet -> Array Int Int -> Array Int IS.IntSet
extendOutput ou f = ou1
  where
    ou1 = listArray (bounds ou) $
          ou ! 0 :
          [ IS.union ous $ ou1 ! fs
          | (ous, fs) <- zip (tail $ elems ou) (elems f) ]

少しとんで6節で、f関数の飛び先をgotoに埋め込んだδを作っている。
これもf関数と同じく、木の低い位置ほど先に計算したい本質がキュー云々という実装で薄まっている。
やはりf関数と同様に、宣言的に定義できる。

buildDelta :: UArray (Int, Char) Int -> Array Int Int -> Array (Int, Char) Int
buildDelta goto f = delta
  where
    (_,ub) = bounds f
    delta = listArray ((0,'a'),(ub,'z')) $ map deltaf $ assocs goto
    deltaf ((s, a), -1) = delta ! (f ! s, a)
    deltaf (_, t) = t

こうして構築されたオートマトンに、入力 $S$ を与えて動かす。
状態 s に到達したとき、output[s] にあるキーワードがその位置を末尾として見つかった。
これをキーワードごとに回数を数えることが問題の要求である。

とりあえずオートマトンの動作の定義アルゴリズム1、ただしδを使う版を素直に書く。

match :: Array (Int, Char) Int -- delta
      -> Array Int IS.IntSet -- output
      -> String -- text string
      -> [IS.IntSet] -- 答えN個、順にその位置で見つかったパターンの集合
match delta output x = snd $ mapAccumL step 0 x
  where
    step :: Int -> Char -> (Int, IS.IntSet)
    step s ai = (t, output ! t)
      where
        t = delta ! (s, ai)

ここで、状態 s に到達するたびに、output[s] の要素それぞれをカウントアップする、とやると手間が掛かりすぎる。
$Q = 5 \times 10^5$ 個のキーワードが全て同じ1文字で、$S$ も全て同じその文字のとき、カウントアップが $NQ = 25 \times 10^{10}$ 回必要になる。$T_i$ に重複はないとは一言も書いてないのだ。

なので、この問題に特化させて、オートマトンの実行では状態ごとの遷移回数だけ数えておき、終わってから、outout の値をその回数倍して足し合わせることにする。

matchCount :: Array (Int, Char) Int -- delta
           -> Array Int IS.IntSet -- output
           -> Int -- パターンの個数Q
           -> String -- text string
           -> [Int] -- 答えQ個、順にパターンの発見回数
matchCount delta output q x = elems pcnt
  where
    scnt, pcnt :: UArray Int Int
    scnt = accumArray (+) 0 (bounds output) [(s, 1) | s <- tail $ scanl (curry (delta !)) 0 x]
    pcnt = accumArray (+) 0 (1,q) [(p, c) | (s,c) <- assocs scnt, p <- IS.elems $ output ! s]

命令型配列の使用

output, f, delta の配列は、木の低い方から順に（上の goto の実装では添え字の昇順でよし）解消できることはわかっているが、相互依存する値を解決する遅延配列で定義したため、その知識を使えず、膨大な量のサンクが発生するために遅い（のだと思う）。

このコードのテイストを極力残して Array.ST で書いた版は、惜しくも TLEx1となった。
ここで、delta は goto と別の配列として作る必要はなく、上書き更新してよいことも盛り込んである。

これをさらに Ix を被せた Vector.Mutable に差し替えたところ、AC, 2730ms, 329MB とぎりぎり間に合った。

これは多分 Data.Sequence が遅いのだろう。
この場合キューの要素は最大で Q 個とわかっているので、その大きさのリングバッファを使えば速くできるだろうけど、面倒なのでやめ。
もう一つ、キーワードを全て一度に扱い、先頭文字で accumArray で選り分けているところも重いだろう。
命令型配列を使うなら、キーワードは一つずつ登録する普通の形にして問題ない。
しかしそうすると、添え字順がノードの高さと揃わなくなるので、以降の計算を木の低い方から実行するための制御が必要になる。
ここでキューを使うと本末転倒。
ノードを高さごとに選り分けた配列を構築することを考える。添え字の上限はキーワードの長さの最大値。
キーワードを登録する際に、新規にノードを割り当てるたびに、それがキーワードの何文字目か、つまりノードの高さで配列に登録する。
このアプローチによる版はAC, 2235ms, 235MB と、余裕のあるタイムだった。

ABC268ExもAho-Corasick法を使う問題らしい。

Suffix Array

公式解説などを眺めると、Suffix Array を使う方が主流のようだ。
それが理解できたらこの問題に戻ってこよう。