Naive Code
クエリ毎に調べる文字列の個数が $N \leq 10^4$
クエリの総数が $M \leq 10^4$
文字列とパターンの両方の長さ上限が $\leq 100$
愚直な方法だと、文字の比較を $10^4 \times 10^4 \times 10^2 = 10^{10}$ 回することになって時間が掛かりすぎるという設定。
解説冒頭にもそう書いてあるのだが、やってみる。
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のデータ構造を定義する。
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な言語では、全てのクエリを束にして、先頭文字で区分けしていくことで一気に木を組み立てる必要がある。
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に先頭から文字列を走らせて、そのとき、到達、通過した全てのノードに貼り付けられたクエリ番号の和集合を求める。
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の要素であるクエリ番号に対して、それぞれ自分の値段を配って足し合わせる、とやればよい。
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)$ ではないのでもっとかかるはずだ。比較でなくハッシュでする辞書なら話は別だが。
まぁそれが想定解だというならやってみよう。
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のまま計算
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]をメモリに広げたりはしていないと思うのだけど。
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のまま計算
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 |
想定解が最速になりました。失礼しました。誤差の範囲のような気もするけど。