103 lines
2.4 KiB
Go
103 lines
2.4 KiB
Go
// Package trie is an implementation of a trie (prefix tree) data structure over byte slices. It provides a
|
|
// small and simple API for usage as a set as well as a 'Node' API for walking the trie.
|
|
package trie
|
|
|
|
// A Trie is a a prefix tree.
|
|
type Trie struct {
|
|
root *Node
|
|
}
|
|
|
|
// New construct a new, empty Trie ready for use.
|
|
func New() *Trie {
|
|
return &Trie{
|
|
root: &Node{},
|
|
}
|
|
}
|
|
|
|
// Insert puts b into the Trie. It returns true if the element was not previously in t.
|
|
func (t *Trie) Insert(b []byte) bool {
|
|
n := t.root
|
|
for _, c := range b {
|
|
next, ok := n.Walk(c)
|
|
if !ok {
|
|
next = &Node{}
|
|
n.branches[c] = next
|
|
n.hasChildren = true
|
|
}
|
|
n = next
|
|
}
|
|
if n.terminal {
|
|
return false
|
|
}
|
|
n.terminal = true
|
|
return true
|
|
}
|
|
|
|
// Contains checks t for membership of b.
|
|
func (t *Trie) Contains(b []byte) bool {
|
|
n := t.root
|
|
for _, c := range b {
|
|
next, ok := n.Walk(c)
|
|
if !ok {
|
|
return false
|
|
}
|
|
n = next
|
|
}
|
|
return n.terminal
|
|
}
|
|
|
|
// PrefixIndex walks through `b` until a prefix is found (terminal node) or it is exhausted.
|
|
func (t *Trie) PrefixIndex(b []byte) int {
|
|
var idx int
|
|
n := t.root
|
|
for _, c := range b {
|
|
next, ok := n.Walk(c)
|
|
if !ok {
|
|
return -1
|
|
}
|
|
if next.terminal {
|
|
return idx
|
|
}
|
|
n = next
|
|
idx++
|
|
}
|
|
if !n.terminal {
|
|
idx = -1
|
|
}
|
|
return idx
|
|
}
|
|
|
|
// Root returns the root node of a Trie. A valid Trie (i.e., constructed with New), always has a non-nil root
|
|
// node.
|
|
func (t *Trie) Root() *Node {
|
|
return t.root
|
|
}
|
|
|
|
// A Node represents a logical vertex in the trie structure.
|
|
type Node struct {
|
|
branches [256]*Node
|
|
terminal bool
|
|
hasChildren bool
|
|
}
|
|
|
|
// Walk returns the node reached along edge c, if one exists. The ok value indicates whether such a node
|
|
// exist.
|
|
func (n *Node) Walk(c byte) (next *Node, ok bool) {
|
|
next = n.branches[int(c)]
|
|
return next, (next != nil)
|
|
}
|
|
|
|
// Terminal indicates whether n is terminal in the trie (that is, whether the path from the root to n
|
|
// represents an element in the set). For instance, if the root node is terminal, then []byte{} is in the
|
|
// trie.
|
|
func (n *Node) Terminal() bool {
|
|
return n.terminal
|
|
}
|
|
|
|
// Leaf indicates whether n is a leaf node in the trie (that is, whether it has children). A leaf node must be
|
|
// terminal (else it would not exist). Logically, if n is a leaf node then the []byte represented by the path
|
|
// from the root to n is not a proper prefix of any element of the trie.
|
|
func (n *Node) Leaf() bool {
|
|
return !n.hasChildren
|
|
}
|