golang 字典树 搜索+sugg 示例

来源:互联网 发布:网络四十大禁书 编辑:程序博客网 时间:2024/05/18 00:01
package mainimport (    "fmt"    "sort"    "sync")type KeyWordKV map[int64]stringtype CharBeginKV map[string][]*KeyWordTreeNodetype PairList []Pairfunc (p PairList) Len() int           { return len(p) }func (p PairList) Less(i, j int) bool { return p[i].V > p[j].V }func (p PairList) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }type Pair struct {    K int64    V int64}type KeyWordTreeNode struct {    // 1, 2, 3    KeyWordIds map[int64]bool    // 百,度, 一, 下    Char string    // 父节点    ParentKeyWordTreeNode *KeyWordTreeNode    // 子节点集合    SubKeyWordTreeNodes map[string]*KeyWordTreeNode}func NewKeyWordTreeNode() *KeyWordTreeNode {    return &KeyWordTreeNode{        KeyWordIds:            make(map[int64]bool, 0),        Char:                  "",        ParentKeyWordTreeNode: nil,        SubKeyWordTreeNodes:   make(map[string]*KeyWordTreeNode, 0),    }}func NewKeyWordTreeNodeWithParams(ch string, parent *KeyWordTreeNode) *KeyWordTreeNode {    return &KeyWordTreeNode{        KeyWordIds:            make(map[int64]bool, 0),        Char:                  ch,        ParentKeyWordTreeNode: parent,        SubKeyWordTreeNodes:   make(map[string]*KeyWordTreeNode, 0),    }}type KeyWordServer struct {    root          *KeyWordTreeNode    kv            KeyWordKV    char_begin_kv CharBeginKV    rw            sync.RWMutex}func NewKeyWordServer() *KeyWordServer {    return &KeyWordServer{        root:          NewKeyWordTreeNode(),        kv:            KeyWordKV{},        char_begin_kv: CharBeginKV{},    }}func (s *KeyWordServer) Put(id int64, keyword string) {    s.rw.Lock()    s.kv[id] = keyword    key_word_tmp_pt := s.root    for _, v := range keyword {        ch := string(v)        if key_word_tmp_pt.SubKeyWordTreeNodes[ch] == nil {            // root:{百:{}}            node := NewKeyWordTreeNodeWithParams(ch, key_word_tmp_pt)            key_word_tmp_pt.SubKeyWordTreeNodes[ch] = node            // 记住不同位置的"百"的起点            s.char_begin_kv[ch] = append(s.char_begin_kv[ch], node)        }        // {百:{度:{}}}        key_word_tree_node := key_word_tmp_pt.SubKeyWordTreeNodes[ch]        key_word_tree_node.KeyWordIds[id] = true        // 更新指针        key_word_tmp_pt = key_word_tmp_pt.SubKeyWordTreeNodes[ch]    }    s.rw.Unlock()}func (s *KeyWordServer) Sugg(keyword string, limit int) []string {    s.rw.RLock()    key_word_tmp_pt := s.root    is_end := true    for _, v := range keyword {        ch := string(v)        if key_word_tmp_pt.SubKeyWordTreeNodes[ch] == nil {            is_end = false            break        }        // 更新指针        key_word_tmp_pt = key_word_tmp_pt.SubKeyWordTreeNodes[ch]    }    if is_end {        // 前缀        ret := make([]string, 0)        ids := key_word_tmp_pt.KeyWordIds        for id, _ := range ids {            ret = append(ret, s.kv[id])            limit --            if limit == 0 {                break            }        }        return ret    }    s.rw.RUnlock()    return make([]string, 0)}func (s *KeyWordServer) Search(keyword string, limit int) []string {    s.rw.RLock()    ids := make(map[int64]int64, 0)    for pos, v := range keyword {        ch := string(v)        begins := s.char_begin_kv[ch]        for _, begin := range begins {            key_word_tmp_pt := begin            next_pos := pos + 1            for len(key_word_tmp_pt.SubKeyWordTreeNodes) > 0 && next_pos < len(keyword) {                // 最大匹配                next_ch := string(keyword[next_pos])                if key_word_tmp_pt.SubKeyWordTreeNodes[next_ch] == nil {                    break                }                key_word_tmp_pt = key_word_tmp_pt.SubKeyWordTreeNodes[next_ch]                next_pos++            }            // 保存结果            for id, _ := range key_word_tmp_pt.KeyWordIds {                ids[id] = ids[id] + 1            }        }    }    // 排序输出果    list := PairList{}    for id, count := range ids {        list = append(list, Pair{            K: id,            V: count,        })    }    if !sort.IsSorted(list) {        sort.Sort(list)    }    if len(list) > limit {        list = list[:limit]    }    ret := make([]string, 0)    for _, item := range list {        ret = append(ret, s.kv[item.K])    }    s.rw.RUnlock()    return ret}func (s *KeyWordServer) DebugPrint() {    fmt.Println("s.kv =", s.kv)    key_word_tmp_pt := s.root    dfs(key_word_tmp_pt)}func dfs(root *KeyWordTreeNode) {    if root == nil {        return    } else {        fmt.Println("s.root =", root.Char)        fmt.Println("s.KeyWordIds =", root.KeyWordIds)        for _, v := range root.SubKeyWordTreeNodes {            dfs(v)        }    }}func main() {    s := NewKeyWordServer()    s.Put(1, "ba")    s.Put(2, "abd")    s.Put(3, "acd")    //s.DebugPrint()    //fmt.Println(s.Sugg("b", 2))    //fmt.Println(s.Search("a", 2))    //fmt.Println(s.Search("b", 2))    //fmt.Println(s.Search("c", 2))    fmt.Println(s.Search("ba", 2))}
原创粉丝点击