Trie implementation

Author: rnbguy

data_structures/Trie/trie.go

@@ -0,0 +1,46 @@
+// Package trie provides Trie data structures in golang.
+//
+// Wikipedia: https://en.wikipedia.org/wiki/Trie
+package trie
+
+// Node represents each node in Trie.
+type Node struct {
+	children map[rune]*Node // map children nodes
+	isLeaf   bool           // current node value
+}
+
+// NewNode creates a new Trie node with initialized
+// children map.
+func NewNode() *Node {
+	n := &Node{}
+	n.children = make(map[rune]*Node)
+	n.isLeaf = false
+	return n
+}
+
+// Insert inserts words at a Trie node.
+func (n *Node) Insert(s string) {
+	curr := n
+	for _, c := range s {
+		next, ok := curr.children[c]
+		if !ok {
+			next = NewNode()
+			curr.children[c] = next
+		}
+		curr = next
+	}
+	curr.isLeaf = true
+}
+
+// Find finds words at a Trie node.
+func (n *Node) Find(s string) bool {
+	curr := n
+	for _, c := range s {
+		next, ok := curr.children[c]
+		if !ok {
+			return false
+		}
+		curr = next
+	}
+	return true
+}

data_structures/Trie/trie_test.go

@@ -0,0 +1,70 @@
+package trie
+
+import (
+	"testing"
+)
+
+func TestTrie(t *testing.T) {
+	n := NewNode()
+
+	insertWords := [...]string{
+		"nikola",
+		"tesla",
+	}
+	checkWords := map[string]bool{
+		"thomas": false,
+		"edison": false,
+		"nikola": true,
+	}
+
+	for _, w := range insertWords {
+		n.insert(w)
+		t.Logf(
+			"added \"%s\" to the Trie.",
+			w,
+		)
+	}
+
+	for k, v := range checkWords {
+		ok := n.find(k)
+		if ok != v {
+			t.Fatalf(
+				"\"%s\" is supposed to be %sin the Trie.",
+				k,
+				map[bool]string{true: "", false: "NOT "}[v],
+			)
+		}
+		t.Logf(
+			"\"%s\" is %sin the Trie.",
+			k,
+			map[bool]string{true: "", false: "NOT "}[ok],
+		)
+	}
+}
+
+func BenchmarkTrie(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		n := NewNode()
+
+		n.insert("nikola")
+		n.insert("tesla")
+
+		n.find("thomas")
+		n.find("edison")
+		n.find("nikola")
+	}
+}
+
+func ExampleNode() {
+	// creates a new node
+	node := NewNode()
+
+	// adds words
+	node.insert("nikola")
+	node.insert("tesla")
+
+	// finds words
+	node.find("thomas") // false
+	node.find("edison") // false
+	node.find("nikola") // true
+}