Posted at

Goで高速な大量prefix match


やりたいこと


たくさんの文字列群prefixesがあります。ある文字列sが、prefixesのいずれかから始まる文字列であるかを判定する関数を書きなさい。もしマッチした場合は、マッチしたprefixも返しなさい。


prefixes := []string{

"aaa",
"bbb",
"ccc",
// ...
}

// こういう感じで動いてほしい
prefix, matched = MatchPrefixes("aaaa", prefixes) // "aaa", true
prefix, matched = MatchPrefixes("aaabbb", prefixes) // "aaa", true
prefix, matched = MatchPrefixes("zz", prefixes) // "", false

素朴にやると、 strings.HasPrefix で順番に見ていくことになります。しかし今回のように大量に一気に、という趣旨だったらもっと高速化できるはずなので、考えてみます。


順番にHasPrefix

比較対象の素朴実装はこれ。線形なアルゴリズムなので、prefixesが長ければ長いほど遅くなるはず。

package main

import "strings"

func LinearPrefix(s string, prefixes []string) (string, bool) {
for _, p := range prefixes {
if strings.HasPrefix(s, p) {
return p, true
}
}
return "", false
}


prefix長さ → prefix → struct{} の二重マップを作る

もしもprefixとかじゃなくて、完全一致であれば、マップを使って O(1) で探索できます。

dictionary := map[string]struct{}{

"aaa": {},
"bbb": {},
"ccc": {},
//...
}

if _, ok := dictionary[s]; ok {
// 完全一致でmatchした
}

prefixマッチだと長さが一定しないのでこの方法が使えませんが、prefixesを長さごとに分類すれば、長さごとにマッチさせることができます。

この方法はprefixesの性質によって速さが変化します。prefixesの中に同じ長さの文字列が多ければ多いほど効率が良くなり、prefixesが全部別の長さだったら線形検索と変わらなくなってしまいます。

package main

type LenMap map[int]map[string]struct{}

func NewLenMap() LenMap {
return make(LenMap)
}

func (m LenMap) Add(p string) {
l := len(p)
_, ok := m[l]
if !ok {
m[l] = make(map[string]struct{}, 1)
}
m[l][p] = struct{}{}
}

func (m LenMap) Match(s string) (string, bool) {
for l, mm := range m {
if len(s) < l {
continue
}
prefix := s[0:l]
if _, ok := mm[prefix]; ok {
return prefix, true
}
}
return "", false
}


Trieを構築

まあ、こういう系統のアルゴリズムならトライ木を作ってしまうのが高速なんじゃないか、ってことで適当に作ってみます。

byteで扱ってもいいのですが、runeの方が書きやすかったのでruneごとにしました。(この方が深さが減るのでbyte単位で作るより速くなるかも?)

package main

type Trie struct {
elm string
children map[rune]*Trie
}

func (t *Trie) Add(s string) {
cur := t
for _, r := range s {
if cur.children == nil {
cur.children = make(map[rune]*Trie, 1)
}

tt, ok := cur.children[r]
if !ok {
tt = new(Trie)
cur.children[r] = tt
}

if tt.elm != "" {
return // NOTE: aaaを登録後aaaaを登録しようとすると、単に無視する
}
cur = tt
}
cur.elm = s
}

func (t *Trie) Match(s string) (string, bool) {
if len(s) == 0 {
return "", true
}

cur := t
for _, r := range s {
if len(cur.children) == 0 {
return cur.elm, true
}

tt, ok := cur.children[r]
if !ok {
return "", false
}
cur = tt
}
return "", false
}


benchmark


prefixesの長さが全部10

prefixesの数=20

BenchmarkLinearPrefix-4          3000000               494 ns/op

BenchmarkLenMapPrefix-4 3000000 473 ns/op
BenchmarkTriePrefix-4 3000000 424 ns/op

prefixesの数=2000

BenchmarkLinearPrefix-4           200000              7977 ns/op

BenchmarkLenMapPrefix-4 3000000 491 ns/op
BenchmarkTriePrefix-4 3000000 493 ns/op

prefixesの数=200000

BenchmarkLinearPrefix-4             2000            984142 ns/op

BenchmarkLenMapPrefix-4 2000000 625 ns/op
BenchmarkTriePrefix-4 1000000 1391 ns/op

prefixの長さが揃っている場合は、やっぱりLenMapの戦略が速そう


prefixesの長さが10~200でランダム

※文字数を増やしたので全体的に時間がかかるようになっています

prefixesの数=20

BenchmarkLinearPrefix-4           200000              5177 ns/op

BenchmarkLenMapPrefix-4 300000 5430 ns/op
BenchmarkTriePrefix-4 300000 5206 ns/op

prefixesの数=2000

BenchmarkLinearPrefix-4           100000             20302 ns/op

BenchmarkLenMapPrefix-4 100000 15640 ns/op
BenchmarkTriePrefix-4 300000 5282 ns/op

prefixesの数=200000

BenchmarkLinearPrefix-4              200           5921473 ns/op

BenchmarkLenMapPrefix-4 30000 40852 ns/op
BenchmarkTriePrefix-4 300000 6331 ns/op

最後のベンチマーク、TriePrefixだけ時間がかかったので、おそらくtrieの構築には相当の時間がかかっていると思われます。(benchの時間に入れていません)


結論

Trieがやはり優秀