Prefix Tree: Trie

Posted on Edited on In

Purpose of this data structure

can search words by prefix in time complexity, O(n).

Recursion Implement

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package main

import (
	"fmt"
)

type Node struct {
	Keys  map[string]*Node
	IsEnd bool
}

func newNode() *Node {
	return &Node{make(map[string]*Node), false}
}

type Trie struct {
	Root *Node
}

func newTrie() *Trie {
	return &Trie{newNode()}
}

func (t *Trie) Add(str string) {
	Add(str, t.Root)
}

func Add(str string, node *Node) {
	if len(str) == 0 {
		node.IsEnd = true
		return
	} else if _, ok := node.Keys[str[:1]]; ok {
		Add(str[1:], node.Keys[str[:1]])
	} else {
		// create new node for this string
		node.Keys[str[:1]] = newNode()
		Add(str[1:], node.Keys[str[:1]])
	}
}

func (t *Trie) HasWord(str string) bool {
	return HasWord(str, t.Root)
}

func HasWord(str string, node *Node) bool {
	if node.IsEnd && len(str) == 0 {
		return true
	} else if len(str) == 0 {
		return false
	}
	if _, ok := node.Keys[str[:1]]; ok {
		return HasWord(str[1:], node.Keys[str[:1]])
	} else {
		return false
	}
}

func (t *Trie) Words() []string {
	words := []string{}
	Print("", t.Root, &words)
	return words
}

func Print(curr string, node *Node, words *[]string) {
	if node.IsEnd {
		*words = append(*words, curr)
	}
	for key := range node.Keys {
		Print(curr+key, node.Keys[key], words)
	}
}

func main() {
	t := newTrie()
	t.Add("apple")
	t.Add("cat")
	t.Add("category")
	t.Add("")
	fmt.Println(t.Words())
	fmt.Println(t.HasWord("apple"))
	fmt.Println(t.HasWord("none"))
	fmt.Println(t.HasWord("category"))
	fmt.Println(t.HasWord("cat"))
	fmt.Println(t.HasWord(""))
}

output

1
2
3
4
5
6
[ apple cat category]
true
false
true
true
true

Forloop Implement

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
type TrieNode struct {
	Children [26]*TrieNode
	IsEnd    bool
}

type Trie struct {
	Root *TrieNode
}


func Constructor() Trie {
    return Trie{&TrieNode{}}
}


func (this *Trie) Insert(word string) {
	curr := this.Root
	for i := range word {
		// check if has word[i] as a child
		k := word[i] - 'a'
		if curr.Children[k] == nil {
			curr.Children[k] = &TrieNode{}
		}
		curr = curr.Children[k]
	}
	curr.IsEnd = true
}

func (this *Trie) Search(word string) bool {
	curr := this.Root
	for i := range word {
		// check if has word[i] as a child
		k := word[i] - 'a'
		if curr.Children[k] == nil {
			return false
		}
		curr = curr.Children[k]
	}
	return curr.IsEnd
}

func (this *Trie) StartsWith(word string) bool {
	curr := this.Root
	for i := range word {
		// check if has word[i] as a child
		k := word[i] - 'a'
		if curr.Children[k] == nil {
			return false
		}
		curr = curr.Children[k]
	}
	return true
}