trie.go 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. /*
  2. Copyright 2016 The Kubernetes Authors.
  3. Licensed under the Apache License, Version 2.0 (the "License");
  4. you may not use this file except in compliance with the License.
  5. You may obtain a copy of the License at
  6. http://www.apache.org/licenses/LICENSE-2.0
  7. Unless required by applicable law or agreed to in writing, software
  8. distributed under the License is distributed on an "AS IS" BASIS,
  9. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. See the License for the specific language governing permissions and
  11. limitations under the License.
  12. */
  13. package util
  14. // A simple trie implementation with Add and HasPrefix methods only.
  15. type Trie struct {
  16. children map[byte]*Trie
  17. wordTail bool
  18. word string
  19. }
  20. // NewTrie creates a Trie and add all strings in the provided list to it.
  21. func NewTrie(list []string) Trie {
  22. ret := Trie{
  23. children: make(map[byte]*Trie),
  24. wordTail: false,
  25. }
  26. for _, v := range list {
  27. ret.Add(v)
  28. }
  29. return ret
  30. }
  31. // Add adds a word to this trie
  32. func (t *Trie) Add(v string) {
  33. root := t
  34. for _, b := range []byte(v) {
  35. child, exists := root.children[b]
  36. if !exists {
  37. child = &Trie{
  38. children: make(map[byte]*Trie),
  39. wordTail: false,
  40. }
  41. root.children[b] = child
  42. }
  43. root = child
  44. }
  45. root.wordTail = true
  46. root.word = v
  47. }
  48. // HasPrefix returns true of v has any of the prefixes stored in this trie.
  49. func (t *Trie) HasPrefix(v string) bool {
  50. _, has := t.GetPrefix(v)
  51. return has
  52. }
  53. // GetPrefix is like HasPrefix but return the prefix in case of match or empty string otherwise.
  54. func (t *Trie) GetPrefix(v string) (string, bool) {
  55. root := t
  56. if root.wordTail {
  57. return root.word, true
  58. }
  59. for _, b := range []byte(v) {
  60. child, exists := root.children[b]
  61. if !exists {
  62. return "", false
  63. }
  64. if child.wordTail {
  65. return child.word, true
  66. }
  67. root = child
  68. }
  69. return "", false
  70. }