Skip to content

Instantly share code, notes, and snippets.

@unnisworld
Last active April 10, 2019 14:36
Show Gist options
  • Save unnisworld/ad7607be45160d0ca2e2ba1dc3db1e93 to your computer and use it in GitHub Desktop.
Save unnisworld/ad7607be45160d0ca2e2ba1dc3db1e93 to your computer and use it in GitHub Desktop.
package dsa;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
public class Trie {
private TrieNode root;
public Trie() {
root = new TrieNode(' ', null);
}
/**
* Inserts a word into the Trie.
*
* @param word
*/
public void insert(String word) {
Map<Character, TrieNode> children = root.childern;
// Keep track of the parent node while we traverse through the trie.
// Parent node info is required while creating a new TrieNode.
TrieNode parentNode = root;
for(int i=0; i< word.length(); i++) {
char currentChar = word.charAt(i);
// Traverse through the Trie to add missing nodes for the word
// that is getting inserted.
TrieNode currentNode;
if (children.containsKey(currentChar)) {
currentNode = children.get(currentChar);
} else {
currentNode = new TrieNode(currentChar, parentNode);
children.put(currentChar, currentNode);
}
// Make the current node as parent node and process it's children.
parentNode = currentNode;
children = currentNode.childern;
// If the current character is the last character,
// mark the current node as a valid Word in the Trie
if (i == word.length()-1) {
currentNode.isEndOfWord = true;
}
}
}
/**
* Displays all words present in a Trie.
*/
public void displayWords(TrieNode currentNode, String str) {
if (currentNode.isEndOfWord) {
System.out.println(str);
}
Map<Character, TrieNode> children = currentNode.childern;
for(Entry<Character, TrieNode> entry : children.entrySet()) {
displayWords(entry.getValue(), str + entry.getKey());
}
}
/**
* Returns a List of suggestions based on the prefix.
*
* @param prefix
* @return List containing words that starts with prefix.
*/
public List<String> suggest(String prefix) {
List<String> words = new ArrayList<String>();
// Locate the prefix in the Trie
TrieNode t = root;
boolean prefixFound = true;
for(int i=0; i< prefix.length(); i++) {
char c = prefix.charAt(i);
t = t.childern.get(c);
if (t == null) {
prefixFound = false;
break;
}
}
// Collect all child words for the prefix.
if (prefixFound) {
collectChildWords(t, words);
Collections.sort(words);
}
return words;
}
private void collectChildWords(TrieNode currentNode, List<String> words) {
if (currentNode.isEndOfWord) {
words.add(currentNode.toString());
}
Map<Character, TrieNode> children = currentNode.childern;
for(Entry<Character, TrieNode> entry : children.entrySet()) {
collectChildWords(entry.getValue(), words);
}
}
public static void main(String[] args) {
Trie trie = new Trie();
trie.insert("dog");
trie.insert("damp");
trie.insert("dance");
trie.insert("damage");
trie.insert("cat");
trie.insert("camera");
trie.insert("catch");
trie.insert("camp");
trie.displayWords(trie.root, "");
List<String> suggestions = trie.suggest("ca");
System.out.println(suggestions);
}
static class TrieNode {
char c;
Map<Character, TrieNode> childern = new HashMap<>();
boolean isEndOfWord;
// Having a parent pointer helps us to implement toString()
TrieNode parent;
TrieNode(char c, TrieNode parent) {
this.c = c;
this.parent = parent;
}
public String toString() {
if (parent == null) {
return "";
}
return parent.toString() + c;
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment