You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

trie.go 2.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. // Package trie is an implementation of a trie (prefix tree) data structure over byte slices. It provides a
  2. // small and simple API for usage as a set as well as a 'Node' API for walking the trie.
  3. package trie
  4. // A Trie is a a prefix tree.
  5. type Trie struct {
  6. root *Node
  7. }
  8. // New construct a new, empty Trie ready for use.
  9. func New() *Trie {
  10. return &Trie{
  11. root: &Node{},
  12. }
  13. }
  14. // Insert puts b into the Trie. It returns true if the element was not previously in t.
  15. func (t *Trie) Insert(b []byte) bool {
  16. n := t.root
  17. for _, c := range b {
  18. next, ok := n.Walk(c)
  19. if !ok {
  20. next = &Node{}
  21. n.branches[c] = next
  22. n.hasChildren = true
  23. }
  24. n = next
  25. }
  26. if n.terminal {
  27. return false
  28. }
  29. n.terminal = true
  30. return true
  31. }
  32. // Contains checks t for membership of b.
  33. func (t *Trie) Contains(b []byte) bool {
  34. n := t.root
  35. for _, c := range b {
  36. next, ok := n.Walk(c)
  37. if !ok {
  38. return false
  39. }
  40. n = next
  41. }
  42. return n.terminal
  43. }
  44. // PrefixIndex walks through `b` until a prefix is found (terminal node) or it is exhausted.
  45. func (t *Trie) PrefixIndex(b []byte) int {
  46. var idx int
  47. n := t.root
  48. for _, c := range b {
  49. next, ok := n.Walk(c)
  50. if !ok {
  51. return -1
  52. }
  53. if next.terminal {
  54. return idx
  55. }
  56. n = next
  57. idx++
  58. }
  59. if !n.terminal {
  60. idx = -1
  61. }
  62. return idx
  63. }
  64. // Root returns the root node of a Trie. A valid Trie (i.e., constructed with New), always has a non-nil root
  65. // node.
  66. func (t *Trie) Root() *Node {
  67. return t.root
  68. }
  69. // A Node represents a logical vertex in the trie structure.
  70. type Node struct {
  71. branches [256]*Node
  72. terminal bool
  73. hasChildren bool
  74. }
  75. // Walk returns the node reached along edge c, if one exists. The ok value indicates whether such a node
  76. // exist.
  77. func (n *Node) Walk(c byte) (next *Node, ok bool) {
  78. next = n.branches[int(c)]
  79. return next, (next != nil)
  80. }
  81. // Terminal indicates whether n is terminal in the trie (that is, whether the path from the root to n
  82. // represents an element in the set). For instance, if the root node is terminal, then []byte{} is in the
  83. // trie.
  84. func (n *Node) Terminal() bool {
  85. return n.terminal
  86. }
  87. // Leaf indicates whether n is a leaf node in the trie (that is, whether it has children). A leaf node must be
  88. // terminal (else it would not exist). Logically, if n is a leaf node then the []byte represented by the path
  89. // from the root to n is not a proper prefix of any element of the trie.
  90. func (n *Node) Leaf() bool {
  91. return !n.hasChildren
  92. }