208.implement-trie-prefix-tree.go 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. type Node struct {
  2. IsKey bool
  3. Children *[26]Node
  4. }
  5. type Trie struct {
  6. Root *Node
  7. }
  8. /** Initialize your data structure here. */
  9. func Constructor() Trie {
  10. return Trie{&Node{}}
  11. }
  12. /** Inserts a word into the trie. */
  13. func (this *Trie) Insert(word string) {
  14. curr := this.Root
  15. var i int
  16. for i = range word {
  17. if curr.Children == nil {
  18. curr.Children = &[26]Node{}
  19. }
  20. curr = &curr.Children[int(word[i]-'a')]
  21. }
  22. curr.IsKey = true
  23. }
  24. /** Returns if the word is in the trie. */
  25. func (this *Trie) Search(word string) bool {
  26. curr := this.Root
  27. var i int
  28. for i = range word {
  29. if curr.Children == nil {
  30. return false
  31. }
  32. curr = &curr.Children[int(word[i]-'a')]
  33. }
  34. return curr.IsKey
  35. }
  36. /** Returns if there is any word in the trie that starts with the given prefix. */
  37. func (this *Trie) StartsWith(prefix string) bool {
  38. curr := this.Root
  39. for i := range prefix {
  40. if curr.Children == nil {
  41. return false
  42. }
  43. curr = &curr.Children[int(prefix[i]-'a')]
  44. }
  45. return curr.IsKey || curr.Children != nil
  46. }
  47. /**
  48. * Your Trie object will be instantiated and called as such:
  49. * obj := Constructor();
  50. * obj.Insert(word);
  51. * param_2 := obj.Search(word);
  52. * param_3 := obj.StartsWith(prefix);
  53. */