Trie

假设我们要实现一个类似于有道词典这样的字典，开始时使用哈希表来存储键值对，后来在实现前缀搜索，模糊匹配等相关功能时，发现使用哈希表作为底层存储的数据结构有一个致命的缺点，基于哈希表的前缀搜索，模糊匹配实现难度非常之大，所以决定换一种数据结构来存储数据。下面我们来介绍一下 Trie 树， Trie 树也叫字典树，它非常适合我们上面说的场景，作为一种树形结构，还经常被用于统计，排序和保存大量的字符串等场景。

data structure

Trie 是一棵R叉树，可以用于保存大量的字符串，每棵树都有一个根节点（ root ）。非叶节点存储一个 rune ，除了指向父节点的指针 parent 之外，并还包含一系列子节点 children ，使用 map 来存储。叶节点标志着一个字符串的结束，其 term 标记为 true ，叶节点中还可以存储一些元数据（ meta ），来表示该字符串的含义。从根节点开始一直到叶节点就可以形成一个完整的字符串。

type Node struct {
    val      rune
    term     bool
    depth    int
    meta     interface{}
    parent   *Node
    children map[rune]*Node
}

type Trie struct {
    root *Node
    size int
}

func New() *Trie {
    return &Trie{
        root: &Node{depth: 0, children: make(map[rune]*Node)}
    }
}

helper function

为了方便后面的搜索操作，我们需要实现两个 helper function ，第一个就是计算一个字符串的掩码，例如对 apple 计算掩码时，分别对’a’，’p’，’p’，’l’，’e’ 这几个 rune 进行位运算，最后返回一个 uint64 ，该掩码在搜索中起着非常重要的作用，因为每个节点都会存储一个掩码（ mask ），我们可以利用该掩码判断该节点及后续字节点存储了什么字符。findNode 方法帮助我们搜索 runes 字符数组中最后一个字符在树中对应的节点。

func maskruneslice(rs []rune) uint64 {
    var m uint64
    for _, r := range rs {
        m |= uint64(1) << uint64(r-'a')
    }
    return m
}

func findNode(node *Node, runes []rune) *Node {
    if node == nil {
        return nil
    }

    if len(runes) == 0 {
        return node
    }

    n, ok := node.Children()[runes[0]]
    if !ok {
        return nil
    }

    var nrunes []rune
    if len(runes) > 1 {
        nrunes = runes[1:]
    } else {
        nrunes = runes[0:0]
    }

    return findNode(n, nrunes)
}

Basic Operation

Trie 树的基本操作有三个，Add ，Find 和 Del 。

Add 方法将 key 字符串加入到树中。对于每一个要加入树中的字符串中的 rune ， 我们构造一个节点，如果相应的 rune 已经存在，此时只需要更新其 mask 值，如果该 rune 不存在，则构建一个中间节点 node = node.NewChild(r, bitmask, nil, false) ，在所有中间节点都构建完毕之后，我们构建一个叶节点结束 Add 操作，叶节点的 val 值设为 nul ，term 为 true ，并存储相应的元数据 meta 。

const nul = 0x0

func (t *Trie) Add(key string, meta interface) *Node {
    t.size++
    runes := []rune(key)
    bitmask := maskruneslice(runes)
    node := t.root
    node.mask |= bitmask
    for i := range runes {
        r := runes[i]
        bitmask := maskruneslice(runes[i:])
        if n, ok := node.children[r]; ok {
            node = n
            node.mask |= bitmask
        } else {
            node = node.NewChild(r, bitmask, nil, false)
        }
    }
    node = node.NewChild(nul, 0, meta, true)
    return node
}

func (n *Node) NewChild(val rune, bitmask uint64, meta interface{}, term bool) *Node {
    node := &Node{
        val: val,
        bitmask: bitmask,
        term: term,
        meta: meta,
        parent: n,
        children: make(map[rune]*Node),
        depth: n.depth + 1,
    }
    n.children[val] = node
    n.mask |= bitmask
    return node
}

Find 方法判断一个特定的字符串是否出现在该树中，我们先通过 findNode 找到相应的 node ，Find 方法用于查找一个完整的字符串，例如我们要查找 app ，但树中只有 apple 这个字符串，所以 findNode 调用时应该返回说不存在 app 这个字符串。当查找 app 这个字符串时，findNode 方法会返回最后一个 p 对应的节点，但这并不说明 app 存在于树中，我们还需要判断该节点是否为叶节点 node.term == true ，如果不是叶节点，则返回 nil 。

func (t *Trie) Find(key string) (*Node, bool) {
    node := findNode(t.Root(), []rune(key))
    if node == nil {
        return nil, false
    }
    node, ok := node.Children()[nul]
    if !ok || !node.term {
        return nil, false
    }
    return node, true
}

Remove 方法同样通过 findNode 方法找到相应的节点，然后从该节点向根节点的方向依次删除相应的 rune ，主要目的是为了更新每个中间节点的mask 值。

func (t *Trie) Remove(key string) {
    var (
        i int
        rs = []rune(key)
        node = findNode(t.Root(), []rune(key))
    )

    t.size--
    for n := node.Parent(); p != nil; n = n.Parent() {
        i++
        if len(n.Children()) > 1 {
            r := rs[len(rs)-i]
            n.RemoveChild(r)
            break
        }
    }
}

func (n *Node) RemoveChild(r rune) {
    delete(n.children, r)
    for nd := n.parent; nd != nil; nd = nd.parent {
        nd.mask ^= nd.mask
        nd.mask |= uint64(1) << uint64(nd.val - 'a')
        for _, c := range nd.children {
            nd.mask |= c.mask
        }
    }
}

Collect

有时候我们搜索单词的时候只记得该单词的前缀，例如我们要搜索 append 这个单词，但是我们只记得 app 这个前缀，我们可以使用 collect 方法帮我们收集所有 app 开头的单词，换句话说就是特定节点下面所有可能的叶节点的集合，开始收集时我们构造一个队列，该队列中仅有 node 一个节点，然后使用广度优先搜索的方式（层序遍历），首先将队列中的第一个节点出队，然后将该节点所有的子节点加入该队列中，如果节点为叶节点，则从该叶节点开始一直向上遍历到根节点的下一层节点（ p.depth == 1 ），直到完成所有单词的收集为止。

func collect(node *Node) []string {
    var (
        keys []string
        n    *Node
        i    int
    )
    nodes := []*Node{node}
    for l := len(nodes); l != 0; l = len(nodes) {
        i := l - 1
        n = nodes[i]
        nodes = nodes[:i]
        for _, c := range n.children {
            nodes = append(nodes, c)
        }
        if n.term {
            word := ""
            for p := n.parent; p.depth != 0; p = p.parent {
                word = string(p.val) + word
            }
            keys = append(keys, word)
        }
    }
    return keys
}

Fuzzy Collect

除了比较准确的前缀匹配之外，有时候我们可能还需要模糊匹配，再举 append 这个单词为例子，我们可能只记得这个单词包含了 apd 这样的模式，使用 Trie 树我们可以非常轻松地实现模糊匹配功能，fuzzycollect 同样利用层序遍历，在每一个循环开始时，对队列中的第一个节点进行出队操作，如果相应的掩码不包含在正在遍历的节点中 (p.node.mask & m) != m ，说明其后面的节点中不包含所有 partial 中包含的 rune 值，直接进入下一个循环 ，如果 p.idx == len(partial) ，说明我们找到了与 partial 相匹配的一个单词，此时进行 collect 操作 。最后同样需要将该节点的子节点加入队列中以进行下一轮的层序遍历。

type potentialSubtree struct {
    idx int
    node *Node
}

func fuzzycollect(node *Node, partial []rune) []string {
    var (
        m    uint64
        i    int
        p    potentialSubtree
        keys []string
    )

    potential := []potentialSubtree{potentialSubtree{0, node}}
    for l := len(potential); l > 0; l = len(potential) {
        i := l - 1
        p = potential[i]
        potential = potential[:i]
        m = maskruneslice(partial[p.idx:])
        if (p.node.mask & m) != m {
            continue
        }

        if p.node.val == partial[p.idx] {
            p.idx++
            if p.idx == len(partial) {
                keys = append(keys, collect(p.node)...)
                continue
            }
        }

        for _, c := range p.node.children {
            potential = append(potential, potentialSubtree{node: c, idx: p.idx})
        }
    }
    return keys
}

search

type ByKeys []string

func (a ByKeys) Len() int           { return len(a) }
func (a ByKeys) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKeys) Less(i, j int) bool { return len(a[i]) < len(a[j]) }

func (t *Trie) HasKeysWithPrefix(key string) bool {
    node := findNode(t.Root(), []rune(key))
    return node != nil
}

func (t Trie) FuzzySearch(pre string) []string {
    keys := fuzzycollect(t.Root(), []rune(pre))
    sort.Sort(ByKeys(keys))
    return keys
}

func (t Trie) PrefixSearch(pre string) []string {
    node := findNode(t.Root(), []rune(pre))
    if node == nil {
        return nil
    }
    return collect(node)
}

func (t *Trie) Keys() []string {
    return t.PrefixSearch("")
}