1package main2 3import (4 "fmt"5 "sort"6 "sync"7 "time"8)9 10type entry struct {11 Text string12 Score int13}14 15type trieNode struct {16 children map[rune]*trieNode17 items []*entry18 mu sync.RWMutex19}20 21func newTrieNode() *trieNode {22 return &trieNode{23 children: make(map[rune]*trieNode),24 items: make([]*entry, 0),25 mu: sync.RWMutex{},26 }27}28 29type PrefixTree struct {30 root *trieNode31 index map[string]*entry32}33 34func newPrefixTree() *PrefixTree {35 return &PrefixTree{36 root: newTrieNode(),37 index: make(map[string]*entry),38 }39}40 41func (p *PrefixTree) Add(text string, score int) {42 cur := p.root43 44 it, ok := p.index[text]45 if !ok {46 it = &entry{47 Text: text,48 Score: score,49 }50 p.index[text] = it51 }52 53 for _, ch := range text {54 cur.mu.RLock()55 _, ok = cur.children[ch]56 57 if !ok {58 cur.mu.RUnlock()59 cur.mu.Lock()60 cur.children[ch] = newTrieNode()61 cur.mu.Unlock()62 } else {63 cur.mu.RUnlock()64 }65 66 cur = cur.children[ch]67 cur.items = append(cur.items, it)68 }69}70 71func (p *PrefixTree) Lookup(prefix string) []string {72 cur := p.root73 var ok bool74 75 for _, ch := range prefix {76 cur, ok = cur.children[ch]77 if !ok {78 return []string{}79 }80 }81 82 items := cur.items83 sort.Slice(items, func(i, j int) bool {84 return items[i].Score > items[j].Score85 })86 87 out := make([]string, 0, len(items))88 for _, it := range items {89 out = append(out, it.Text)90 }91 92 return out93}94 95func main() {96 tree := newPrefixTree()97 go tree.Add("monitor", 100)98 go tree.Add("module", 30)99 go tree.Add("magnet", 20)100 go tree.Add("rocket", 10)101 102 time.Sleep(1 * time.Second)103 104 fmt.Println(tree.Lookup("m")) // ["monitor", "module", "magnet"]105 fmt.Println(tree.Lookup("mo")) // ["monitor", "module"]106 fmt.Println(tree.Lookup("ma")) // ["magnet"]107 fmt.Println(tree.Lookup("r")) // ["rocket"]108}