paiza「文字列収集 (paizaランク S 相当)」をHaskellで

[問題](https://paiza.jp/works/mondai/real_event/word_collection)

# Naive Code

クエリ毎に調べる文字列の個数が $N \leq 10^4$
クエリの総数が $M \leq 10^4$
文字列とパターンの両方の長さ上限が $\leq 100$

愚直な方法だと、文字の比較を $10^4 \times 10^4 \times 10^2 = 10^{10}$ 回することになって時間が掛かりすぎるという設定。
解説冒頭にもそう書いてあるのだが、やってみる。

```haskell
import Control.Monad
import Data.List

main :: IO ()
main = do
  [n,m] <- map read . words <$> getLine
  sps <- replicateM n $ do
    [s,pp] <- words <$> getLine
    return (s, read pp)
  qs <- replicateM m getLine
  let ans = solve n m sps qs
  mapM_ print ans

solve :: Int -> Int -> [(String, Int)] -> [String] -> [Int]
solve n m sps qs = [sum [p | (s,p) <- sps, isPrefixOf q s] | q <- qs]
```

`ByteString` すら使わないナメくさりコード。実質1行。

結果：
> 0.40 秒 スコア 100 点 ランク S 相当

あれぇ？
出題当時の計算機の性能から向上した恩恵？

# Trie

ランクS問題なのでもっと真面目にやろう。
とりあえず、Trieのデータ構造を定義する。

```haskell
import Data.Char
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

data Trie = Node IS.IntSet (IM.IntMap Trie)
```

`IntMap`はキーが`Int`限定なので、`Data.Char.ord` を使って文字コードでインデックスする。
`IntSet`は、Trieのそのノードに到達したときに一致したクエリの番号の集合を持たせる。

mutableな計算なら、クエリを一つずつTrieに追加するのが自然なコードだが、immutableな言語では、全てのクエリを束にして、先頭文字で区分けしていくことで一気に木を組み立てる必要がある。

```haskell
import Data.Array

buildTrie :: [(Int, [Char])] -> Trie
buildTrie iqs = Node is im
  where
    arr = accumArray (flip (:)) [] ('a','z') [(x, (i, xs)) | (i, x:xs) <- iqs]
    is = IS.fromList [i | (i, "") <- iqs]
    im = IM.fromDistinctAscList [(ord c, buildTrie ixs) | (c, ixs) <- assocs arr, not (null ixs)]
```
引数には、クエリ番号とクエリ文字列（のtail）の組のリストを与える。
それらの登録されたTrieを再帰的に構築する。
`arr` で、先頭の文字ごとに分岐先を仕分けする。
`is` はこのノードで完了するクエリの番号を集める。
`im` は、`arr`の中で要素が入った文字について、その内容で再帰的にTrieを作り、できたものをまとめて `IntMap` にして Trie の枝分かれを作っている。

このTrieに先頭から文字列を走らせて、そのとき、到達、通過した全てのノードに貼り付けられたクエリ番号の和集合を求める。

```haskell
runTrie :: Trie -> String -> IS.IntSet
runTrie trie = loop trie IS.empty . map ord
  where
    loop (Node is im) js (c:cs)
      | IM.member c im = loop (im IM.! c) (IS.union is js) cs
    loop (Node is _) js _ = IS.union is js
```

辿る先があるかぎり `loop` の再帰で先に進む。文字列が完了するか、枝がないか、とにかく進む先がなくなれば終了する。どちらにせよ到達したノードの `IntSet` の内容は集めていく。

ここまで準備ができれば、問題を解くには、クエリからTrieを作り、それぞれの文字列について`runTrie`を行い、結果の`IntSet`の要素であるクエリ番号に対して、それぞれ自分の値段を配って足し合わせる、とやればよい。

```haskell
solve :: Int -> Int -> [(String, Int)] -> [String] -> [Int]
solve n m sps qs = elems accs
  where
    trie = buildTrie $ zip [1..] qs
    accs = accumArray (+) 0 (1, m) [(j, p) | (s,p) <- sps, j <- IS.elems $ runTrie trie s]
```

結果：
> 0.06 秒 スコア 100 点 ランク S 相当

実際のところ、最初にこれをサクっと書いて提出して、いくら効率的になったといっても速すぎると思って愚直コードを書いたら通ってしまったのでこれを書くことにしたのだけど、これで王道制覇。一件落着、か？

# 想定解

勝ち誇って公式の解説を開いたら、全然違う解き方をしている。あれぇ？

「文字列の接頭辞になっているクエリを見つけたい」という要求に答えるのに、
「文字列の接頭辞すべてをキー、その値段を値として合計を登録した巨大な辞書を用意しておけば、クエリについてその辞書を検索するだけ」という発想はなかった。
辞書の要素数が文字列の長さの総和（$L$とする）になり、登録は $O(L \log L)$ 検索は $O(M \log L)$ となるように見えるが、文字列の辞書順比較は $O(1)$ ではないのでもっとかかるはずだ。比較でなくハッシュでする辞書なら話は別だが。

まぁそれが想定解だというならやってみよう。

```haskell
solve :: Int -> Int -> [(String, Int)] -> [String] -> [Int]
solve n m sps qs = [M.findWithDefault 0 q m | q <- qs]
  where
    m = M.fromListWith (+) [(s1, p) | (s,p) <- sps, s1 <- tail $ inits s]
```

実質2行。

結果：
> 0.18 秒 スコア 100 点 ランク S 相当

Trieによる方法より遅いけど、2行で実装が済むし愚直解法よりは速いので、アリですかね？

# 終わりに

スキルチェックでHaskell使えるようにしてホスイ…

# 追記：ByteString版

ここまでやったらByteStringで全開にした版もやっておかないと、と気づいたので追加。

愚直版：ByteStringのまま計算
```haskell
solveNaive :: Int -> Int -> [(BS.ByteString, Int)] -> [BS.ByteString] -> [Int]
solveNaive n m sps qs = [sum [p | (s,p) <- sps, BS.isPrefixOf q s] | q <- qs]
```

Trie版：`BS.unpack`で`String`に戻すだけ
といっても多分ByteStringから1文字ずつ遅延評価で持ってくるから、`[Char]`をメモリに広げたりはしていないと思うのだけど。
```haskell
solveTrie :: Int -> Int -> [(BS.ByteString, Int)] -> [BS.ByteString] -> [Int]
solveTrie n m sps qs = elems accs
  where
    trie = buildTrie $ zip [1..] $ map BS.unpack qs
    accs = accumArray (+) 0 (1, m) [(j, p) | (s,p) <- sps, j <- IS.elems $ runTrie trie $ BS.unpack s]
```

Map版：ByteStringのまま計算
```haskell
solve :: Int -> Int -> [(BS.ByteString, Int)] -> [BS.ByteString] -> [Int]
solve n m sps qs = [M.findWithDefault 0 q m | q <- qs]
  where
    m = M.fromListWith (+) [(s1, p) | (s,p) <- sps, s1 <- tail $ BS.inits s]
```

|   | タイム(sec) |
|:-:|:-:|
| 愚直 | 0.04 |
| Trie | 0.05 |
| Map  | 0.03 |

想定解が最速になりました。失礼しました。誤差の範囲のような気もするけど。