Trie Implementation in Java

Question

I am trying to implement a Trie class in Java which just supports insertion and searching operation up until now. I think it works perfectly on the several examples I have tried but I am not able to come up with a tricky test case to check if it functions perfectly or not. Can you take a look at my implementation and see if it's the right way to code up a trie data structure?

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;

/* This class represents the structure of node that will be stored in a trie */
class TrieNode implements Comparable<TrieNode>, Comparator<TrieNode>{

    // the alphabet to be stored in a node
    Character letter;
    // this variable stores if the node is a terminating point of a valid word in a dictionary
    boolean isLeaf;
    // this variable stores the count of the words that end at this node
    int count;
    // this set stores the set of child nodes attached to this node
    ArrayList<TrieNode> childNodes;

    // this is the parameterless constructor for the class
    public TrieNode(){

    }

    // this is the constructor for the TrieNode class
    public TrieNode(Character letter, boolean isLeaf, int count, ArrayList<TrieNode> childNodes){
        this.letter = letter;
        this.isLeaf = isLeaf;
        this.count = count;
        this.childNodes = childNodes;
    }

    // this method is used for sorting all of the Trie nodes according to their alphabet value
    public int compareTo(TrieNode o) {
        return this.letter.compareTo(o.letter);
    }

    @Override
    public int compare(TrieNode o1, TrieNode o2) {
        return o1.compareTo(o2);
    }
}

/* This class represents the Trie data structure and is used for inserting and searching in a Trie */
class Trie{

    // this reference variable stores the root node of a Trie data structure
    TrieNode root;

    // this is the constructor for the Trie class
    public Trie(){
        this.root = new TrieNode(' ', false, 0, new ArrayList<TrieNode>());
    }

    // this function is used to insert a character into a Trie
    public void insertTrie(String word){

        ArrayList<TrieNode> temp = this.root.childNodes;
        Collections.sort(temp, new TrieNode());

        for(int i = 0; i < word.length(); ++i){
            boolean found = false;
            for(int j = 0; j < temp.size(); ++j){
                if(temp.get(j).letter == word.charAt(i)){
                    temp = temp.get(j).childNodes;
                    Collections.sort(temp, new TrieNode());
                    found = true;
                    break;
                }
            }

            if(found)
                continue;

            TrieNode newNode;

            if(i+1 == word.length()){
                newNode = new TrieNode(word.charAt(i), true, 0, new ArrayList<TrieNode>());
            }else{
                newNode = new TrieNode(word.charAt(i), false, 0, new ArrayList<TrieNode>());
            }

            temp.add(newNode);
            Collections.sort(temp, new TrieNode());
            temp = newNode.childNodes;
            Collections.sort(temp, new TrieNode());
        }
    }

    // this function returns true if the word being searched for exists in the trie
    public boolean searchTrie(String word){

        ArrayList<TrieNode> temp = this.root.childNodes;

        for(int i = 0; i < word.length() - 1; ++i){
            boolean found = false;
            for(int j = 0; j < temp.size(); ++j){
                if(temp.get(j).letter != word.charAt(i))
                    continue;
                else{
                    found = true;
                    temp = temp.get(j).childNodes;
                    break;
                }
            }

            if(!found)
                return false;
        }

        for(int i = 0; i < temp.size(); ++i){
            if(word.charAt(word.length() - 1) == temp.get(i).letter && temp.get(i).isLeaf)
                return true;
        }

        return false;
    }
}

/* for the sake of simplicity we consider only lowercase english alphabets */
public class DataStructure1_TrieOperations {
    public static void main(String[] args){

        String text = "bell bid stock spock sell shell";
        String[] dictionary = text.split("\\s+");

        Trie tr1 = new Trie();
        for(int i = 0; i < dictionary.length; ++i){
            tr1.insertTrie(dictionary[i]);
        }

        System.out.println(tr1.searchTrie("shell"));
    }
}

Answer 1

    // this variable stores if the node is a terminating point of a valid word in a dictionary
    boolean isLeaf;
    // this variable stores the count of the words that end at this node
    int count;

Why have both? You can get the stated meaning for isLeaf with count == 0. Why keep an extra variable that does the same thing? Also, a leaf is normally a node with no children. A trie node can terminate a word and still have children.

    ArrayList<TrieNode> childNodes;

As a general rule, you should define variables of the type of the interface rather than the implementation. In other words, you'd normally say

    List<TrieNode> childNodes;

You also may want to consider the advantages of using a Map rather than a List. In particular, if you use a SortedMap, then you don't need to sort childNodes before using them. The insertion will automatically maintain sorted order.

        for(int i = 0; i < word.length(); ++i){

Try to avoid calling a function on each iteration that always returns the same value.

        for ( int i = 0, n = word.length() - 1; i < n; ++i ) { 

This way it calls the function once and saves the value.

        Collections.sort(temp, new TrieNode());

The Comparator form is unnecessary here. You can just say

    Collections.sort(temp);

Because TrieNode is a Comparable type.

            if(found)
                continue;

Try to avoid the single statement form of loop control structures. It's harder to quickly scan and follow plus it can lead to a class of bug that can be nasty to debug. If always using the block form saves even one debug session, you'll come out ahead in the long term.

            TrieNodeAS newNode;

            if(i+1 == word.length()){
                newNode = new TrieNodeAS(word.charAt(i), true, 0, new ArrayList<TrieNodeAS>());
            }else{
                newNode = new TrieNodeAS(word.charAt(i), false, 0, new ArrayList<TrieNodeAS>());
            }

A few things. First, there's a much shorter way to write this:

            boolean isLeaf = i+1 == word.length();
            TrieNodeAS newNode = new TrieNodeAS(word.charAt(i), isLeaf, 0, new ArrayList<TrieNodeAS>());

The entire if/else structure is unnecessary. The expression in the if returns a Boolean value, so pass it directly. You don't even need the intermediate variable, although it does make the logic a little clearer.

Second, this doesn't track count properly. You set it to 0 in both cases even though it should be 1 in the first case. If you're not going to use it, you might as well get rid of it.

Third, why create an empty list and pass it to the constructor? It's a constructor. Make it create its own list.

            for(int j = 0; j < temp.size(); ++j){
                if(temp.get(j).letter == word.charAt(i)){
                    temp = temp.get(j).childNodes;

If you can try to use the foreach form rather than the old C style for loop. In this case, you only use j in temp.get(j) (which you use twice instead of saving the value the first time).

            for ( TrieNode current : temp ) {
                if ( current.letter == word.charAt(i) ) {
                    temp = current.childNodes;

This saves calling size() on each iteration of the loop, plus it eliminates the get(j) calls.

The whole TrieNode.java using a SortedMap rather than an ArrayList:

class TrieNode {
    int count;
    SortedMap<Character, TrieNode> childNodes;

    public TrieNode(int count){
        this.count = count;
        this.childNodes = new TreeMap<Character, TrieNode>();
    }
}

And the accompanying Trie.java:

class Trie{
    TrieNode root;

    public Trie(){
        this.root = new TrieNode(0);
    }

    public void insertTrie(String word) {
        SortedMap<Character, TrieNode> nextLetters = this.root.childNodes;

        // all but the last letter in the word
        for ( int i = 0, n = word.length() - 1; i < n; ++i ) {
            char letter = word.charAt(i);

            if ( ! nextLetters.containsKey(letter) ) {
                nextLetters.put(letter, new TrieNode(0));
            }

            nextLetters = nextLetters.get(letter).childNodes;
        }

        // the last letter, so set or update count appropriately
        char letter = word.charAt(word.length() - 1);
        TrieNode node = nextLetters.get(letter);
        if ( null == node ) {
            nextLetters.put(letter, new TrieNode(1));
        } else {
            node.count++;
        }
    }

    public boolean searchTrie(String word) {
        SortedMap<Character, TrieNode> nextLetters;
        TrieNode node = this.root;

        for ( Character letter : word.toCharArray() ) {
            nextLetters = node.childNodes;
            node = nextLetters.get(letter);

            if ( node == null ) {
                return false;
            }
        }

        return node.count > 0;
    }
}

I tested it with main:

public static void main(String[] args) {
    String[] words = {"bell", "bid", "stock", "spock", "sell", "shell"};

    Trie trie = new Trie();
    for ( String word : words ) {
        trie.insertTrie(word);
    }

    System.out.println(trie.searchTrie("bi"));
    System.out.println(trie.searchTrie("bid"));
    System.out.println(trie.searchTrie("bide"));
    System.out.println(trie.searchTrie("shelf"));
    System.out.println(trie.searchTrie("shell"));
}

Note: I did not find any functional problems with your original code. However, I think that my version is shorter and more robust.