Golang – Insert and Search Trie

Golang – Insert and Search Trie

Problem Statement:

Write a program to insert and search in trie data structure.

We suggest you think about a solution before reading further…

Solution:

A simple representation of trie data structure to store and search english words is

Insert :
  1) For each character in the word, get index.
  2) Check whether character is present in the trie or not
  3) If present, move to the next node
  4) If not, create a new trie node and append with current node.
  5) At the end of word, set endOfWords to true.
Search :
  1) For each character in the given word, get index
  2) check whether character is present in the trie or not
  3) If present, move to next node
  4) If not, return false
  5) At the end of word, check endOfWords flag. If endOfWords == true, returns true else return false.

Code to insert and search in Trie

package main

import (
	"fmt"
)

const ALPHABET_SIZE = 26

type TrieNode struct {
	children   [ALPHABET_SIZE]*TrieNode
	endOfWords bool
}

func getNode() *TrieNode {
	node := &TrieNode{}
	node.endOfWords = false

	for i := 0; i < ALPHABET_SIZE; i++ {
		node.children[i] = nil
	}

	return node
}

func insert(root *TrieNode, key string) {
	temp := root

	for i := 0; i < len(key); i++ {
		index := key[i] - 'a'
		if temp.children[index] == nil {
			temp.children[index] = getNode()
		}
		temp = temp.children[index]
	}

	temp.endOfWords = true
}

func search(root *TrieNode, key string) bool {
	temp := root

	for i := 0; i < len(key); i++ {
		index := key[i] - 'a'
		if temp.children[index] != nil {
			temp = temp.children[index]
		} else {
			return false
		}
	}

	return (temp != nil && temp.endOfWords)
}

func main() {
	words := []string{"a", "and", "an", "go", "golang", "man", "mango"}
	root := getNode()

	for i := 0; i < len(words); i++ {
		insert(root, words[i])
	}

	fmt.Println("contains words [a]: ", search(root, "a"))
	fmt.Println("contains words [mango]: ", search(root, "mango"))
	fmt.Println("contains words [lang]: ", search(root, "lang"))
}

Output

contains words [a]:  true
contains words [mango]:  true
contains words [lang]:  false

Stay tuned for more updates and tutorials

Leave a Reply

Your email address will not be published. Required fields are marked *