Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Prefix tree written by using Trie. Supports find prefixes in O(n.m.k) where k is length of the prefix and O(n.m) comes from Depth First Search
import java.util.ArrayList;
import java.util.Map.Entry;
import java.util.Set;
public class Dictionary {
private TrieNode root;
public Dictionary() {
root = new TrieNode(null);
}
public void insertWord(String word) {
if (!isValid(word)) {
throw new IllegalArgumentException("Word is null or empty");
}
TrieNode child = null;
String key = word.substring(0, 1);
if (!this.root.getChildren().containsKey(key)) {
child = new TrieNode(key);
this.root.getChildren().put(key, child);
} else {
child = this.root.getChildren().get(key);
}
insert(word.substring(1), child);
}
private void insert(String word, TrieNode node) {
TrieNode currentNode = null;
String key = word.substring(0, 1);
if (!node.getChildren().containsKey(key)) {
currentNode = new TrieNode(key);
node.getChildren().put(key, currentNode);
} else {
currentNode = node.getChildren().get(key);
}
if (word.length() == 1) {
currentNode.setTerminal(true);
} else {
insert(word.substring(1), currentNode);
}
}
public boolean searchWord(String word) {
if (!isValid(word)) {
throw new IllegalArgumentException();
}
TrieNode node = null;
String key = word.substring(0, 1);
if (!root.getChildren().containsKey(key)) {
return false;
} else {
node = root.getChildren().get(key);
}
return search(word.substring(1), node);
}
private boolean search(String word, TrieNode node) {
String key = word.substring(0, 1);
if (!node.getChildren().containsKey(key)) {
return false;
}
if (word.length() == 1) {
return true;
} else {
return search(word.substring(1), node.getChildren().get(key));
}
}
public ArrayList<String> findPrefixes(String prefix) {
if (!isValid(prefix)) {
throw new IllegalArgumentException();
}
ArrayList<String> prefixList = new ArrayList<>();
String key = prefix.substring(0, 1);
if (!root.getChildren().containsKey(key)) {
return null;
}
TrieNode endNode = getEndNode(prefix.substring(0), root);
if (endNode != null) {
return getPrefixes(endNode, prefix, prefixList);
}
return null;
}
private ArrayList<String> getPrefixes(TrieNode endNode, String prefix, ArrayList<String> prefixList) {
Set<Entry<String, TrieNode>> entrySet = endNode.getChildren().entrySet();
for (Entry<String, TrieNode> entry : entrySet) {
TrieNode value = entry.getValue();
dfsVisit(value, prefix, prefixList);
}
return prefixList;
}
private void dfsVisit(TrieNode value, String str, ArrayList<String> prefixList) {
str = str + value.getKey();
Set<Entry<String, TrieNode>> entrySet = value.getChildren().entrySet();
for (Entry<String, TrieNode> entry : entrySet) {
TrieNode node = entry.getValue();
dfsVisit(node, str, prefixList);
}
if (value.isTerminal()) {
prefixList.add(str);
}
}
private TrieNode getEndNode(String prefix, TrieNode node) {
String key = prefix.substring(0, 1);
TrieNode child = node.getChildren().get(key);
if (child == null) {
return null;
}
if (prefix.length() == 1) {
return child;
} else {
return getEndNode(prefix.substring(1), child);
}
}
private boolean isValid(String str) {
if (str == null || str.isEmpty()) {
return false;
}
return true;
}
}
import java.util.HashMap;
public class TrieNode {
private final String key;
private boolean isTerminal;
private final HashMap<String, TrieNode> children;
public TrieNode(String key) {
this.key = key;
this.setTerminal(false);
this.children = new HashMap<>();
}
public String getKey() {
return key;
}
public boolean isTerminal() {
return isTerminal;
}
public void setTerminal(boolean isTerminal) {
this.isTerminal = isTerminal;
}
public HashMap<String, TrieNode> getChildren() {
return children;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment