I implemented a Trie data structure in Java, to test my understanding of the concept. I tried (pun intended :) ) to follow TDD steps along the way (i.e., add first a failing test case, and implement functionality after ensuring that it really fails). I'm posting here the end result. Besides the implementation itself, I'm also interested in the quality of the tests.
I'm mostly interested in the following aspects:
- Is this a correct implementation of the Trie data structure? If not, then what are the mistakes?
- Can the performance or the effectiveness (both in speed, as well as in readability) of the implementation be improved?
- Can you think of any test case, which would cause an issue in this implementation, besides the trivial NPE issues, if
null
is passed to the methods (I intentionally left out theif (word == null) ...
checks to make the code shorter)? (In other words, a test that exposes a bug.) - Would you improve anything in the unit tests?
Of course, any other suggestion is also welcome.
What I already know could be improved:
I intentionally left out the
isLeaf
part for now. (Which tells, for each node, if it is the end of a complete word. E.g., if we had a "path" forhello
, then both theo
at the end, as well as the lastl
could have theisLeaf
flag set to true.)string2Coll
method could be written in a more elegant way, using Java 8 (but for some reason my Eclipse can not set Java 8 code generation, although I have JDK for Java 8 installed -- and at this point, I'm not too interested in that aspect)
For the record, I checked in this code into my github repository, and probably I'll improve it based on the answers from here.
Trie.java
package trie;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
public class Trie<E> {
private static class Node<E> {
private Map<E, Node<E>> children = new HashMap<>();
Node<E> addChild(E child) {
children.putIfAbsent(child, new Node<E>());
return children.get(child);
}
Node<E> getChild(E child) {
return children.get(child);
}
Node<E> removeChild(E child) {
return children.remove(child);
}
boolean hasAtMostOneChild() {
return children.size() < 2;
}
}
private Node<E> root = new Node<>();
public boolean contains(Collection<E> word) {
Node<E> node = root;
for (E ch: word) {
Node<E> next = node.getChild(ch);
if (next == null) {
return false;
} else {
node = next;
}
}
return true;
}
public void add(Collection<E> word) {
Node<E> node = root;
for (E ch: word) {
Node<E> next = node.getChild(ch);
if (next == null) {
next = node.addChild(ch);
}
node = next;
}
}
public void remove(Collection<E> word) {
Node<E> node = root;
Node<E> lastFork = root;
E childAtLastFork = word.iterator().hasNext() ? word.iterator().next() : null;
for (E ch: word) {
if (!node.hasAtMostOneChild()) {
lastFork = node;
childAtLastFork = ch;
}
node = node.getChild(ch);
if (node == null) {
break;
}
}
lastFork.removeChild(childAtLastFork);
}
}
TrieTest.java
package trie;
import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertThat;
import java.util.ArrayList;
import java.util.Collection;
import org.junit.Test;
public class TrieTest {
@Test
public void emptyTrieShouldContainEmptyWord() {
Trie<Character> trie = new Trie<>();
assertThat(trie.contains(string2Coll("")), is(true));
}
@Test
public void emptyTrieShouldNotContainWord() {
Trie<Character> trie = new Trie<>();
assertThat(trie.contains(string2Coll("hello")), is(false));
}
@Test
public void addZeroLengthWordShouldNotCauseException() {
Trie<Character> trie = new Trie<>();
trie.add(string2Coll(""));
}
@Test
public void wordShouldBeContainedIfAdded() {
Trie<Character> trie = new Trie<>();
trie.add(string2Coll("hello"));
assertThat(trie.contains(string2Coll("hello")), is(true));
}
@Test
public void wordShouldNotBeContainedIfRemoved() {
Trie<Character> trie = new Trie<>();
trie.add(string2Coll("hello"));
trie.remove(string2Coll("hello"));
assertThat(trie.contains(string2Coll("hello")), is(false));
}
@Test
public void removingNonExistingWordShouldNotCauseException() {
Trie<Character> trie = new Trie<>();
trie.remove(string2Coll("hello"));
assertThat(trie.contains(string2Coll("hello")), is(false));
}
@Test
public void removingOneWordShouldNotRemoveAnother() {
Trie<Character> trie = new Trie<Character>();
trie.add(string2Coll("hello"));
trie.add(string2Coll("help"));
trie.remove(string2Coll("help"));
assertThat(trie.contains(string2Coll("help")), is(false));
assertThat(trie.contains(string2Coll("hello")), is(true));
}
@Test
public void removingZeroLengthWordShouldNotCauseException() {
Trie<Character> trie = new Trie<>();
trie.remove(string2Coll(""));
assertThat(trie.contains(string2Coll("")), is(true));
}
private static Collection<Character> string2Coll(String str) {
Collection<Character> ret = new ArrayList<>(str.length());
for (char ch: str.toCharArray()) {
ret.add(ch);
}
return ret;
}
}