Trie（字典树）的Java实现

简单实现了一个具有CRUD操作能力的Trie。CRUD操作即插入（Create），读取（Read），改变（Update）和删除（Delete）。

删除是基于当前结点的

count

实现的。当

count

为0时代表当前结点应该被完全删除。

Trie的所有操作均为

O(n)

。其中

n

为所要操作的单词的长度。

public KeyNotFound extends Exception {
public KeyNotFound(String message) {
super(message);
}
}

public class Trie {
class TrieNode {
char curr;
TrieNode[] children;
private boolean isEnd;
String data;
private int count;
TrieNode (char curr, boolean isEnd, String data) {
// Cast all the chars to lower case
curr = Character.toLowerCase(curr);
this.curr = curr;
this.isEnd = isEnd;
this.data = data;
this.children = new TrieNode[26];
this.count = 1;
}

void setEnd(boolean end) {
this.isEnd = end;
}

void increaseCount() {
this.count++;
}

int getCount() {
return this.count;
}

void decreaseCount() {
this.count--;
}
}

private TrieNode root;

public Trie() {
/** Initialize the Trie. */
this.root = new TrieNode(' ', false, "");
}

public void insert(String word, String data) {
/** Inserts a word and corresponding data into the trie. */
char[] chars = word.toCharArray();
TrieNode currPos = this.root;
for (int i = 0; i < chars.length; i++) {
int index = chars[i] - 'a';
// If a node corresponding to current substring does not exists, create one.
if (currPos.children[index] == null) {
TrieNode newNode = new TrieNode(chars[i], false, "");
currPos.children[index] = newNode;
}
// If this is the end of the word, store the data in to the Trie and set end.
if (i == chars.length - 1) {
currPos.children[index].setEnd(true);
currPos.children[index].data = data;
}
// Move one level down
currPos = currPos.children[index];
}
}

public boolean search(String word) {
/** Returns if the word is in the Trie. */
char[] chars = word.toCharArray();
TrieNode currPos = this.root;
for (int i = 0; i < chars.length; i++) {
int index = chars[i] - 'a';
// If no match find, return false
if (currPos.children[index] == null) {
return false;
}
if (i == chars.length - 1 && currPos.children[index].isEnd) {
return true;
}
currPos = currPos.children[index];
}
return false;
}

public boolean startsWith(String prefix) {
/** Returns if there is any word in the Trie that starts with the given prefix. */
char[] chars = prefix.toCharArray();
TrieNode currPos = this.root;
for (int i = 0; i < chars.length; i++) {
int index = chars[i] - 'a';
if (currPos.children[index] == null) {
return false;
}
currPos = currPos.children[index];
}
return true;
}

public String getData(String key) throws KeyNotFound {
/**
* Gets the corresponding data for the given key if the key exists in the Trie.
* Throw KeyNotFound error if no such key exists
*/
char[] chars = key.toCharArray();
TrieNode currPos = this.root;
for (int i = 0; i < chars.length; i++) {
int index = chars[i] - 'a';
// If no match find, return empty string
if (currPos.children[index] == null) {
throw new KeyNotFound("Given key does not exists in the Trie!");
}
currPos = currPos.children[index];
}
if (!currPos.isEnd) {
throw new KeyNotFound("Given key does not exists in the Trie!");
}
return currPos.data;
}

public boolean setData(String key, String newData) {
/**
* Change the data of an existing key.
* If changed successfully, return true.
* If the key doesn't exists, return false.
*/
char[] chars = key.toCharArray();
TrieNode currPos = this.root;
for (int i = 0; i < chars.length; i++) {
int index = chars[i] - 'a';
// If no match find, cannot store the data, return false
if (currPos.children[index] == null) {
return false;
}
currPos = currPos.children[index];
}
if (!currPos.isEnd) {
return false;
}
currPos.data = newData;
return true;
}

public void delete(String word) {
/**
* Delete the given word in the Trie if exists
*/
// If the word doesn't exist in the Trie, should delete anything
if (!search(word)) {
return;
}
char[] chars = word.toCharArray();
TrieNode currPos = this.root;
for (int i = 0; i < chars.length; i++) {
int index = chars[i] - 'a';
TrieNode tmp = currPos.children[index];
if (currPos.children[index].getCount() == 1) {
currPos.children[index] = null;
} else {
currPos.children[index].decreaseCount();
}

// if a word is deleted but the prefix still exists, delete the data and flip isEnd flag.
if (i == chars.length - 1) {
if (currPos.children[index] != null) {
currPos.children[index].data = "";
currPos.children[index].setEnd(false);
}
}
currPos = tmp;
}

}
}