Update index.go

Author: masx200

implement-trie-prefix-tree/index.go

@@ -1,24 +1,43 @@
 package implement_trie_prefix_tree
-type Trie struct {
 
+type Trie struct {
+	children [26]*Trie
+	isEnd    bool
 }
 
-
 func Constructor() Trie {
-
+	return Trie{}
 }
 
-
-func (this *Trie) Insert(word string)  {
-
+func (t *Trie) Insert(word string) {
+	node := t
+	for _, ch := range word {
+		ch -= 'a'
+		if node.children[ch] == nil {
+			node.children[ch] = &Trie{}
+		}
+		node = node.children[ch]
+	}
+	node.isEnd = true
 }
 
-
-func (this *Trie) Search(word string) bool {
-
+func (t *Trie) SearchPrefix(prefix string) *Trie {
+	node := t
+	for _, ch := range prefix {
+		ch -= 'a'
+		if node.children[ch] == nil {
+			return nil
+		}
+		node = node.children[ch]
+	}
+	return node
 }
 
+func (t *Trie) Search(word string) bool {
+	node := t.SearchPrefix(word)
+	return node != nil && node.isEnd
+}
 
-func (this *Trie) StartsWith(prefix string) bool {
-
+func (t *Trie) StartsWith(prefix string) bool {
+	return t.SearchPrefix(prefix) != nil
 }