Here is the Java trie implementation, that we can use for many problems.
First thing we need is TrieNode, which represents the individual node structure. The collection of these trienodes can be used to create Trie.
TrieNode Class
When implementing a dictionary, store the corresponding value in final nodes, and null in non-final nodes.
Also, we don’t need to store the char content
anymore, as it is now part of map.
public class TrieNode {
private boolean isLeaf;
private Map<Character, TrieNode> children = new HashMap<>();
}
Sometimes people also use names like isWord
, isEnd
for isLeaf
.
Trie Class
Based on this, we had following Trie class definition:
public class Trie{
private TrieNode root;
public Trie() {
this.root = new TrieNode();
}
public void insert(String word) {
}
public boolean exists(String word) {
}
public boolean startsWith(String prefix) {
}
public List<String> wordsByPrefix(String prefix) {
}
}
Code
Java
Now, we can use this class as static class in Trie.
Full Trie class
public class Trie {
public class TrieNode {
private boolean isLeaf;
private Map<Character, TrieNode> children = new HashMap<>();
}
private TrieNode root;
public Trie() {
root = new TrieNode();
}
// Inserts a word into the trie.
public void insert(String word) {
TrieNode curr = this.root;
for (char c: word.toCharArray()) {
curr.children.putIfAbsent(c, new TrieNode());
curr = curr.children.get(c);
}
curr.isLeaf = true;
}
// Returns if the word is in the trie.
public boolean exists(String word) {
TrieNode curr = root;
for (char c: word.toCharArray()) {
if (!curr.children.containsKey(c)) {
return false;
}
curr = curr.children.get(c);
}
return curr.isLeaf;
}
// Returns if there is any word in the trie
// that starts with the given prefix.
public boolean startsWith(String prefix) {
TrieNode curr = root;
for (char c: prefix.toCharArray()) {
if (!curr.children.containsKey(c)) {
return false;
}
curr = curr.children.get(c);
}
return true;
}
public List<String> wordsByPrefix(String prefix) {
List<String> ans = new ArrayList<>();
TrieNode curr = this.root;
// Navigate to the end of the prefix
for (char c: prefix.toCharArray()) {
if (!curr.children.containsKey(c)) {
// If prefix is not present, return empty list
return ans;
}
curr = curr.children.get(c);
}
// Recursively find all words starting with the given prefix
dfs(curr, new StringBuilder(prefix), ans);
return ans;
}
private void dfs(TrieNode node, StringBuilder path, List<String> ans) {
if (node.isLeaf) {
ans.add(path.toString());
}
for (Map.Entry<Character, TrieNode> entry: node.children.entrySet()) {
// For each child, append the char and search deeper
path.append(entry.getKey());
dfs(entry.getValue(), path, ans);
// Backtrack to previous state to explore other paths
path.deleteCharAt(path.length() - 1);
}
}
// Returns the shortest matching prefix between the string and Trie
// eg. Trie = cat, s = cattle => output = cat
// Trie = cattle123 , s = cattle => cattle
// Trie = cat, s= battle => output = ""
// Trie = cat + cattle, s = cattle => cat
public String shortestMatchingPrefix(String word) {
TrieNode curr = root;
StringBuilder prefixBuilder = new StringBuilder();
for (char ch : word.toCharArray()) {
curr = curr.children.get(ch);
if (curr == null) {
// The word has deviated from any path in the Trie
return "";
}
prefixBuilder.append(ch);
if (curr.isLeaf) {
// The current node marks the end of a word; shortest matching prefix found
return prefixBuilder.toString();
}
}
return "";
}
public void delete(String word) {
delete(this.root, word, 0);
}
private boolean delete(TrieNode current, String word, int index) {
if (index == word.length()) {
// When end of word is reached only delete if current.endOfWord is true.
if (!current.isLeaf) {
return false;
}
current.isLeaf = false;
// If current has no other mapping then return true
return current.children.size() == 0;
}
char ch = word.charAt(index);
TrieNode node = current.children.get(ch);
if (node == null) {
return false;
}
boolean shouldDeleteCurrentNode = delete(node, word, index + 1) && !node.isLeaf;
if (shouldDeleteCurrentNode) {
// If true is returned then delete the mapping of character and trienode reference from map.
current.children.remove(ch);
// Return true if no mappings are left in the map.
return current.children.size() == 0;
}
return false;
}
}