LeetCode 1804. Implement Trie II (Prefix Tree) in Java

In this post I will solve LeetCode 1804. Implement Trie II (Prefix Tree) problem using Java and the VSCode IDE on a Windows computer.

In a previous post we discussed and implemented a solution for LeetCode 208. Implement Trie (Prefix Tree) in Java which is similar yet different to this one.

Please note that I carried away experimenting with the current problem as I did with the previous one. Due to this I implemented two solutions. After the first one was accepted I tried improving performance and readability. That created a second solution. Due to my approach I ended up with a large amount of code which I will post and comment on. For the first implementation I will go light on comments.

Let’s start by reading the requirements for the current problem.

A trie (pronounced as "try") or prefix tree is a tree data structure 
used to efficiently store and retrieve keys in a dataset of strings. 
There are various applications of this data structure, such as autocomplete and spellchecker.

Implement the Trie class:

* Trie()
	Initializes the trie object.

* void insert(String word)
	Inserts the string word into the trie.
	
o int countWordsEqualTo(String word)
	Returns the number of instances of the string word in the trie.
	
o int countWordsStartingWith(String prefix)
	Returns the number of strings in the trie that have the string prefix as a prefix.
	
o void erase(String word) Erases the string word from the trie.

Constraints:

o 1 <= word.length, prefix.length <= 2000
o word and prefix consist only of `lowercase` English letters.
o At most 3 * 10^4 calls in total will be made to 
  insert, countWordsEqualTo, countWordsStartingWith, and erase.
o It is guaranteed that for any function call to erase,
  the string word will exist in the trie.

Related Topics:

o Hash Table
o String
o Design
o Trie

In this problem we need to create a Trie with five methods. Each method comes with a short but precise description. Please take a few moments and think about the specific features and based on them come up with the class members and an approach for the methods.

Once you have an idea on how to count the full words and prefixes, the key is to concentrate on the `insert` and `remove` methods. At least that is how I approached the problem.

One more thing, if interested in the problem, please visit the LeetCode site and make sure you get the current description directly from the source. Problems may slightly change, additional test cases may become available, or additional hints might have been added.

class Trie {

    public Trie() {
        
    }
    
    public void insert(String word) {
        
    }
    
    public int countWordsEqualTo(String word) {
        
    }
    
    public int countWordsStartingWith(String prefix) {
        
    }
    
    public void erase(String word) {
        
    }
}

/**
 * Your Trie object will be instantiated and called as such:
 
 * Trie obj = new Trie();
 * obj.insert(word);
 * int param_2 = obj.countWordsEqualTo(word);
 * int param_3 = obj.countWordsStartingWith(prefix);
 * obj.erase(word);
 */

The Trie class signature seems very adequate for the requirements.

In the comments section the method used by LeetCode to check our solution is described.

I will implement a simple test scaffold so I can solve the problem on my computer. The simplest approach is to develop the solution using the online IDE provided by LeetCode.

Sorry for the length of this post, but like I said earlier, I just got carried away and wrote this Victorian novel.

Trie, insert, insert, countWordsEqualTo, countWordsStartingWith, erase, countWordsEqualTo, countWordsStartingWith, erase, countWordsStartingWith
null, apple,  apple,  apple,             app,                    apple, apple,             app,                    apple, app
main <<<      n: 10
main <<< cmdArr: [Trie, insert, insert, countWordsEqualTo, countWordsStartingWith, erase, countWordsEqualTo, countWordsStartingWith, erase, countWordsStartingWith]
main <<< argArr: [null, apple, apple, apple, app, apple, apple, app, apple, app]
insert <<< word ==>apple<==
insert <<< p: (ch: a prefixCnt: 1 wordCnt: 0 str: a)
insert <<< p: (ch: p prefixCnt: 1 wordCnt: 0 str: ap)
insert <<< p: (ch: p prefixCnt: 1 wordCnt: 0 str: app)
insert <<< p: (ch: l prefixCnt: 1 wordCnt: 0 str: appl)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: apple)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 1 str: apple)
insert <<< word ==>apple<==
insert <<< p: (ch: a prefixCnt: 2 wordCnt: 0 str: a)
insert <<< p: (ch: p prefixCnt: 2 wordCnt: 0 str: ap)
insert <<< p: (ch: p prefixCnt: 2 wordCnt: 0 str: app)
insert <<< p: (ch: l prefixCnt: 2 wordCnt: 0 str: appl)
insert <<< p: (ch: e prefixCnt: 2 wordCnt: 1 str: apple)
insert <<< p: (ch: e prefixCnt: 2 wordCnt: 2 str: apple)
erase <<< word ==>apple<==
erase <<<      p: (ch: * prefixCnt: 0 wordCnt: 0 str: null)
erase <<<   next: (ch: a prefixCnt: 2 wordCnt: 0 str: a)
erase <<< delete: false
erase <<<      p: (ch: a prefixCnt: 1 wordCnt: 0 str: a)
erase <<<   next: (ch: p prefixCnt: 2 wordCnt: 0 str: ap)
erase <<< delete: false
erase <<<      p: (ch: p prefixCnt: 1 wordCnt: 0 str: ap)
erase <<<   next: (ch: p prefixCnt: 2 wordCnt: 0 str: app)
erase <<< delete: false
erase <<<      p: (ch: p prefixCnt: 1 wordCnt: 0 str: app)
erase <<<   next: (ch: l prefixCnt: 2 wordCnt: 0 str: appl)
erase <<< delete: false
erase <<<      p: (ch: l prefixCnt: 1 wordCnt: 0 str: appl)
erase <<<   next: (ch: e prefixCnt: 2 wordCnt: 2 str: apple)
erase <<< delete: false
erase <<< p: (ch: e prefixCnt: 1 wordCnt: 1 str: apple)
erase <<< word ==>apple<==
erase <<<      p: (ch: * prefixCnt: 0 wordCnt: 0 str: null)
erase <<<   next: (ch: a prefixCnt: 1 wordCnt: 0 str: a)
erase <<< delete: true
erase <<<      p: (ch: a prefixCnt: 0 wordCnt: 0 str: a)
erase <<<   next: (ch: p prefixCnt: 1 wordCnt: 0 str: ap)
erase <<< delete: true
erase <<<      p: (ch: p prefixCnt: 0 wordCnt: 0 str: ap)
erase <<<   next: (ch: p prefixCnt: 1 wordCnt: 0 str: app)
erase <<< delete: true
erase <<<      p: (ch: p prefixCnt: 0 wordCnt: 0 str: app)
erase <<<   next: (ch: l prefixCnt: 1 wordCnt: 0 str: appl)
erase <<< delete: true
erase <<<      p: (ch: l prefixCnt: 0 wordCnt: 0 str: appl)
erase <<<   next: (ch: e prefixCnt: 1 wordCnt: 1 str: apple)
erase <<< delete: true
erase <<< p: (ch: e prefixCnt: 0 wordCnt: 0 str: apple)
main <<< output: [null, null, null, 2, 2, null, 1, 1, null, 0]


Trie, insert, insert, insert,    insert, erase, erase, erase,  erase,     insert, countWordsStartingWith, countWordsEqualTo
null, apple,  a,      applebees, append, a,     apple, append, applebees, cake,   c,                      cake
main <<<      n: 12
main <<< cmdArr: [Trie, insert, insert, insert, insert, erase, erase, erase, erase, insert, countWordsStartingWith, countWordsEqualTo]  
main <<< argArr: [null, apple, a, applebees, append, a, apple, append, applebees, cake, c, cake]
insert <<< word ==>apple<==
insert <<< p: (ch: a prefixCnt: 1 wordCnt: 0 str: a)
insert <<< p: (ch: p prefixCnt: 1 wordCnt: 0 str: ap)
insert <<< p: (ch: p prefixCnt: 1 wordCnt: 0 str: app)
insert <<< p: (ch: l prefixCnt: 1 wordCnt: 0 str: appl)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: apple)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 1 str: apple)
insert <<< word ==>a<==
insert <<< p: (ch: a prefixCnt: 2 wordCnt: 0 str: a)
insert <<< p: (ch: a prefixCnt: 2 wordCnt: 1 str: a)
insert <<< word ==>applebees<==
insert <<< p: (ch: a prefixCnt: 3 wordCnt: 1 str: a)
insert <<< p: (ch: p prefixCnt: 2 wordCnt: 0 str: ap)
insert <<< p: (ch: p prefixCnt: 2 wordCnt: 0 str: app)
insert <<< p: (ch: l prefixCnt: 2 wordCnt: 0 str: appl)
insert <<< p: (ch: e prefixCnt: 2 wordCnt: 1 str: apple)
insert <<< p: (ch: b prefixCnt: 1 wordCnt: 0 str: appleb)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: applebe)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: applebee)
insert <<< p: (ch: s prefixCnt: 1 wordCnt: 0 str: applebees)
insert <<< p: (ch: s prefixCnt: 1 wordCnt: 1 str: applebees)
insert <<< word ==>append<==
insert <<< p: (ch: a prefixCnt: 4 wordCnt: 1 str: a)
insert <<< p: (ch: p prefixCnt: 3 wordCnt: 0 str: ap)
insert <<< p: (ch: p prefixCnt: 3 wordCnt: 0 str: app)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: appe)
insert <<< p: (ch: n prefixCnt: 1 wordCnt: 0 str: appen)
insert <<< p: (ch: d prefixCnt: 1 wordCnt: 0 str: append)
insert <<< p: (ch: d prefixCnt: 1 wordCnt: 1 str: append)
erase <<< word ==>a<==
erase <<<      p: (ch: * prefixCnt: 0 wordCnt: 0 str: null)
erase <<<   next: (ch: a prefixCnt: 4 wordCnt: 1 str: a)
erase <<< delete: false
erase <<< p: (ch: a prefixCnt: 3 wordCnt: 0 str: a)
erase <<< word ==>apple<==
erase <<<      p: (ch: * prefixCnt: 0 wordCnt: 0 str: null)
erase <<<   next: (ch: a prefixCnt: 3 wordCnt: 0 str: a)
erase <<< delete: false
erase <<<      p: (ch: a prefixCnt: 2 wordCnt: 0 str: a)
erase <<<   next: (ch: p prefixCnt: 3 wordCnt: 0 str: ap)
erase <<< delete: false
erase <<<      p: (ch: p prefixCnt: 2 wordCnt: 0 str: ap)
erase <<<   next: (ch: p prefixCnt: 3 wordCnt: 0 str: app)
erase <<< delete: false
erase <<<      p: (ch: p prefixCnt: 2 wordCnt: 0 str: app)
erase <<<   next: (ch: l prefixCnt: 2 wordCnt: 0 str: appl)
erase <<< delete: false
erase <<<      p: (ch: l prefixCnt: 1 wordCnt: 0 str: appl)
erase <<<   next: (ch: e prefixCnt: 2 wordCnt: 1 str: apple)
erase <<< delete: false
erase <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: apple)
erase <<< word ==>append<==
erase <<<      p: (ch: * prefixCnt: 0 wordCnt: 0 str: null)
erase <<<   next: (ch: a prefixCnt: 2 wordCnt: 0 str: a)
erase <<< delete: false
erase <<<      p: (ch: a prefixCnt: 1 wordCnt: 0 str: a)
erase <<<   next: (ch: p prefixCnt: 2 wordCnt: 0 str: ap)
erase <<< delete: false
erase <<<      p: (ch: p prefixCnt: 1 wordCnt: 0 str: ap)
erase <<<   next: (ch: p prefixCnt: 2 wordCnt: 0 str: app)
erase <<< delete: false
erase <<<      p: (ch: p prefixCnt: 1 wordCnt: 0 str: app)
erase <<<   next: (ch: e prefixCnt: 1 wordCnt: 0 str: appe)
erase <<< delete: true
erase <<<      p: (ch: e prefixCnt: 0 wordCnt: 0 str: appe)
erase <<<   next: (ch: n prefixCnt: 1 wordCnt: 0 str: appen)
erase <<< delete: true
erase <<<      p: (ch: n prefixCnt: 0 wordCnt: 0 str: appen)
erase <<<   next: (ch: d prefixCnt: 1 wordCnt: 1 str: append)
erase <<< delete: true
erase <<< p: (ch: d prefixCnt: 0 wordCnt: 0 str: append)
erase <<< word ==>applebees<==
erase <<<      p: (ch: * prefixCnt: 0 wordCnt: 0 str: null)
erase <<<   next: (ch: a prefixCnt: 1 wordCnt: 0 str: a)
erase <<< delete: true
erase <<<      p: (ch: a prefixCnt: 0 wordCnt: 0 str: a)
erase <<<   next: (ch: p prefixCnt: 1 wordCnt: 0 str: ap)
erase <<< delete: true
erase <<<      p: (ch: p prefixCnt: 0 wordCnt: 0 str: ap)
erase <<<   next: (ch: p prefixCnt: 1 wordCnt: 0 str: app)
erase <<< delete: true
erase <<<      p: (ch: p prefixCnt: 0 wordCnt: 0 str: app)
erase <<<   next: (ch: l prefixCnt: 1 wordCnt: 0 str: appl)
erase <<< delete: true
erase <<<      p: (ch: l prefixCnt: 0 wordCnt: 0 str: appl)
erase <<<   next: (ch: e prefixCnt: 1 wordCnt: 0 str: apple)
erase <<< delete: true
erase <<<      p: (ch: e prefixCnt: 0 wordCnt: 0 str: apple)
erase <<<   next: (ch: b prefixCnt: 1 wordCnt: 0 str: appleb)
erase <<< delete: true
erase <<<      p: (ch: b prefixCnt: 0 wordCnt: 0 str: appleb)
erase <<<   next: (ch: e prefixCnt: 1 wordCnt: 0 str: applebe)
erase <<< delete: true
erase <<<      p: (ch: e prefixCnt: 0 wordCnt: 0 str: applebe)
erase <<<   next: (ch: e prefixCnt: 1 wordCnt: 0 str: applebee)
erase <<< delete: true
erase <<<      p: (ch: e prefixCnt: 0 wordCnt: 0 str: applebee)
erase <<<   next: (ch: s prefixCnt: 1 wordCnt: 1 str: applebees)
erase <<< delete: true
erase <<< p: (ch: s prefixCnt: 0 wordCnt: 0 str: applebees)
insert <<< word ==>cake<==
insert <<< p: (ch: c prefixCnt: 1 wordCnt: 0 str: c)
insert <<< p: (ch: a prefixCnt: 1 wordCnt: 0 str: ca)
insert <<< p: (ch: k prefixCnt: 1 wordCnt: 0 str: cak)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 0 str: cake)
insert <<< p: (ch: e prefixCnt: 1 wordCnt: 1 str: cake)
main <<< output: [null, null, null, null, null, null, null, null, null, null, 1, 1]

Each test case is separated from others by two blank lines. I had a few test cases which I ended up removing most and leaving only two. The first is the one provided by LettCode. The second I created to make sure additions and deletions would work fine.

Our test code, which is NOT PART OF THE SOLUTION, reads the first two input lines. The first holds the commands in the order they need to be invoked, and the second line contains the argument, if any, that needs to be passed to the method in question. A null implies that an argument is not required for the matching command / method.

Our test code appears to read the two input lines, store the values into two String[] arrays, display the length of the arrays (actually it checks that both arrays have the same length) and then display the necessary values. This is done to make sure that all is well before we start the implementation of the Trie class.

The answer is in the last line of each test case. The text in between contains messages generated by the `insert` and the `erase` methods. Once they are up and running the other Trie class methods should follow nicely. We will cover in some detail such messages when we look at the code for the methods in the second implementation.

    /**
     * Test scaffold
     * !!!! NOT PART OF SOLUTION !!!!
     * @throws IOException
     */
    public static void main(String[] args) throws IOException {
        
        // **** open buffered reader ****
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        // **** read commands (method calls) ****
        String[] cmdArr = Arrays.stream(br.readLine().trim().split(", "))
                                .map(s -> s.trim())
                                .toArray(size -> new String[size]);
        
        // **** get array size ****
        int n = cmdArr.length;
    
        // **** read arguments to commands (method calls) ****
        String[] argArr = Arrays.stream(br.readLine().trim().split(", "))
                                .map(s -> s.trim())
                                .toArray(size -> new String[size]);

        // **** get array size ****
        int m = argArr.length;

        // **** close buffered reader ****
        br.close();

        // **** compare array lengths ****
        if (n != m) {
            System.err.println("main <<< n: " + n + " != m: " + m);
            System.exit(-1);
        }

        // ???? ????
        System.out.println("main <<<      n: " + n);
        System.out.println("main <<< cmdArr: " + Arrays.toString(cmdArr));
        System.out.println("main <<< argArr: " + Arrays.toString(argArr));

        // **** ****
        Trie trie           = null;
        Integer[] output    = new Integer[n];

        // **** process commands in array ****
        for (int i = 0 ; i < n; i++) {

            // **** process this command ****
            switch (cmdArr[i]) {
                case "Trie":
                    trie = new Trie();
                    output[i] = null;
                break;

                case "insert":
                    trie.insert(argArr[i]);
                    output[i] = null;
                break;

                case "countWordsEqualTo":
                    output[i] = trie.countWordsEqualTo(argArr[i]);
                break;

                case "countWordsStartingWith":
                    output[i] = trie.countWordsStartingWith(argArr[i]);
                break;

                case "erase":
                    trie.erase(argArr[i]);
                    output[i] = null;
                break;

                case "anyWords":
                    System.out.println("main <<< trie: " + trie.anyWords(argArr[i]));
                break;

                default:
                    System.err.println("main <<< unexpected cmdArr[" + i + "]: " + cmdArr[i]);
                    System.exit(-1);
                break;
            }

            // ???? ????
            // List<String> words = new ArrayList<>();
            // trie.findWords(trie, words);
            // System.out.println("main <<< anyWords: " + words.toString());
        }

        // **** display output ****
        System.out.println("main <<< output: " + Arrays.toString(output));
    }

The test scaffold is a little more complicated than in other posts, yet it is simple to follow.

It starts by reading the two input lines into the `cmdArr` and `argArr` String[] arrays. Note that in the input lines I added some spaces for the arguments to line up with the associated commands. We need to remove such blank spaces. Such operation is performed while reading the input streams.

We collect the length of the two arrays in the variables `n` and `m`. We then compare them to make sure they match.

After declaring a couple variables we enter a loop in which we call the specified methods using the specified arguments.

Note that the results of the method calls are stored in the `output` Integer[] array. After the loop is done we display the contents of the array so we can compare it to the expected results. Please note that in the first test case the expected results are provided by LeetCode. On the second we need to figure them out.

    /**
     * Trie class
     */
    static class Trie {


        // **** class members ****
        public Trie[]   children;
        public int      prefixCnt;
        public int      wordCnt;

        // ???? experimenting ????
        public char     ch;
		public String   str;
        public boolean  end;
        public Trie     prev;


        /**
         * Constructor.
         */
        public Trie() {
            this.children   = new Trie[26];
            this.prefixCnt  = 0;
            this.wordCnt    = 0;

            // ???? experimenting ????
            this.ch         = '*';
        }


        /**
         * Constructor.
         */
        public Trie(char ch) {
            this.children   = new Trie[26];
            this.prefixCnt  = 0;
            this.wordCnt    = 0;

            // ???? experimenting ????
            this.ch         = ch;
        }


        /**
         * Constructor.
         */
        public Trie(char ch, String str) {
            this.children   = new Trie[26];
            this.prefixCnt  = 0;
            this.wordCnt    = 0;

            // ???? experimenting ????
            this.ch         = ch;
            this.str        = str;
        }


        /**
         * Generate string representation of Trie.
         * Some experimental members are displayed.
         * 
         * !!!! NOT PART OF SOLUTION !!!!
         */
        @Override
        public String toString() {

            // **** initialization ****
            StringBuilder sb = new StringBuilder("(");

            // **** ****
            sb.append("ch: " + this.ch);
            sb.append(" prefixCnt: " + this.prefixCnt);
            sb.append(" wordCnt: " + this.wordCnt);
            sb.append(" str: " + this.str);

            // **** ****
            sb.append(")");

            // **** ****
            return sb.toString(); 
        }
		
		// **** other class methods follow... ****

We need to figure out the class members for the Trie class. After some experimentation the first three are all we need. Note that I experimented with an additional four variables. I believe that all variables were used in the first implementation. As far as I can remember in the second implementation I used less additional values. This was done to provide additional information for the `println` messages.

        /**
         * Inserts the string word into the trie.
         */
        public void insert0(String word) {

            // **** initialization ****
            Trie p = this;

            // **** traverse word inserting characters ****
            for (int i = 0; i < word.length(); i++) {

                // **** for ease of use ****
                char ch = word.charAt(i);

                // **** ****
                Trie prev = p;
    
                // **** create new child  ****
                if (p.children[ch - 'a'] == null)
                    p.children[ch - 'a'] = new Trie(ch, word.substring(0, i + 1));

                // **** point to next child ****
                p = p.children[ch - 'a'];

                // **** ****
                p.prev = prev;
            }

            // **** increment prefixCnt ****
            p.prefixCnt++;

            // **** end of word ****
            p.end = true;
        }

This is the implementation of the `insert` method for the first pass. I will cover the first implementation now. Please note that the output you see in the test cases was for the second implementation; NOT THE FIRST!

The idea is to traverse the `word` string one character at a time and be able to use an existing node in the trie, or add nodes as needed. Of key importance is the update of the `prefixCnt` and the `end` variables in the nodes.

        /**
         * Erases the string word from the trie.
         * It is guaranteed that for any function call to erase, 
         * the string word will exist in the trie.
         */
        public void erase0(String word) {
            
            // **** initialization ****
            Trie p = this;

            // **** traverse the word one character at a time ****
            for (char ch : word.toCharArray()) {

                // **** point to next node ****
                p = p.children[ch - 'a'];

                // **** update current node ****
                if (p.str.equals(word)) {
                    
                    // **** decrement counter ****
                    p.prefixCnt--;

                    // **** flag that word is not in the Trie ****
                    if (p.prefixCnt == 0) p.end = false;
                }
            }
        }

The `erase` method is used to erase a specific word from the trie.

We start at the root and go down traversing the proper nodes. Note that the `prefixCnt` and the `end` variables are updated as needed.

Of interest is that the erased words are not removed from the trie. At some point I was removing them, but I decided to leave them since chances are that the words are valid in the English dictionary, so sooner or later they would be reused. Of course, in production, such an arbitrary decision should have been approved by the development team.

        /**
         * Returns the number of instances of the string word in the trie.
         */
        public int countWordsEqualTo0(String word) {

            // **** initialization ****
            Trie p = this;

            // **** traverse word ****
            for (char ch : word.toCharArray()) {

                // **** character not found ****
                if (p.children[ch - 'a'] == null) return 0;

                // **** continue search ****
                p = p.children[ch - 'a'];
            }

            // **** check if this is not a word ****
            if (p.end == false) return 0;

            // **** return prefixCnt of words equal to word ****
            return p.prefixCnt;
        }

In this method we return the number of times the `word` string is in the trie. The code seems to be self explanatory.

        /**
         * Returns the number of strings in the trie that have 
         * the string prefix as a prefix.
         */
        public int countWordsStartingWith0(String prefix) {

            // **** initialization ****
            Trie p = this;

            // **** traverse the prefix one charracter at a time ****
            for (char ch : prefix.toCharArray()) {

                // **** prefix not found ****
                if (p.children[ch - 'a'] == null) return 0;

                // **** move to next child ****
                p = p.children[ch - 'a'];
            }

            // **** return the number of words that start with prefix ****
            return p.prefixCnt;
        }

This method is used to return the number of times the specified `prefix` is in the trie.

Note that the answer would be in the last node of the trie reached by `p` when processing the last character held in the `prefix` argument.

        /**
         * Return count of words that start here.
         * Recursive method.
         *
         * !!!! NOT PART OF SOLUTION !!!!
         */
        private int countEnds(Trie p) {

            // **** base case ****
            if (p == null) return 0;

            // **** initialization ****
            int wordCnt = 0;

            // **** recurse ****
            for (int i = 0; i < 26; i++)
                wordCnt += countEnds(p.children[i]);

            // **** ****
            if (p.end == true) wordCnt += p.wordCnt;

            // **** return total number of words in the trie ****
            return wordCnt;
        }

In this recursive method IS NOT PART OF THE SOLUTION. It is used to count the total numbers of words in the trie. It is a recursive function. I changed the operation of the function while experimenting. At this time it should probably be renamed to `wordCount`.

        /**
         * Returns a list of words if there are one or more 
         * previously inserted string words that start with
         * the specified `prefix`.
         * If no words are found returns an empty list.
         * 
         * !!!! NOT PART OF THE SOLUTION !!!!
         */
        public List<String> anyWords(String prefix) {
        
            // **** initialization ****
            List<String> words = new ArrayList<>();
        
            // **** sanity checks ****
            if (prefix.length() == 0 || prefix.equals(" ")) return words;
        
            // **** to traverse trie ****
            Trie p = this;
        
            // **** get to the end of the prefix ****
            for (int i = 0; prefix.charAt(0) != '*' && i < prefix.length(); i++) {
        
                // **** get current character ****
                char ch = prefix.charAt(i);
        
                // **** ****
                if (p.children[ch - 'a'] == null) return words;
        
                // **** continue search ****
                p = p.children[ch - 'a'];
            }
        
            // **** search for full words starting here ****
            findWords(p, words);
        
            // **** return list of possible words ****
            return words;
        }

This is another function that I used to experiment. The name does no longer match the implementation. It seems to look for completed words, not for prefixes. This is why the function is flagged as NOT PART OF THE SOLUTION.

        /**
         * Find full words in the trie starting at p
         * and add them to the list.
         * Recursive call.
         * 
         * !!!! NOT PART OF THE SOLUTION !!!!
         */
        public void findWords(Trie p, List<String>words) {
        
            // **** end condition ****
            if (p == null) return;
        
            // **** traverse trie ****
            for (Trie child : p.children)
                findWords(child, words);

            // **** add full world to list ****
            if (p.wordCnt != 0) words.add(p.str + " (" + p.wordCnt + ")");
        }

This recursive function seems to generate a list of words and associated counts in the trie. It seems that it should work. All full words are flagged by the `wordCnt` as being > 0. During development I changed names of variables as needed in the Trie class. In all implementations it counted the end of a word. That said; this function is NOT PART OF THE SOLUTION.

Now let’s get to the final implementation!!!

Please note that flagging the nodes with the `ch` and the `str` variables is NOT PART OF THE SOLUTION. I populated them to better explain what is going on.

        /**
         * Inserts the string word into the trie.
         */
        public void insert(String word) {

            // ???? ????
            System.out.println("insert <<< word ==>" + word + "<==");

            // **** initialization ****
            Trie p = this;

            // **** traverse the word one character at a time ****
            for (int i = 0; i < word.length(); i++) {

                // **** for ease of use ****
                char ch = word.charAt(i);
                int ndx = ch - 'a';

                // **** create and append new Trie node ****
                if (p.children[ndx] == null)
                    p.children[ndx] = new Trie(ch, word.substring(0, i + 1));

                // **** move to next child ****
                p = p.children[ndx];

                // **** increment prefix count *****
                p.prefixCnt++;

                // ???? ????
                System.out.println("insert <<< p: " + p);
            }

            // **** increment word count ****
            p.wordCnt++;

            // ???? ????
            System.out.println("insert <<< p: " + p);
        }

The idea is to traverse the `word` and if needed,  add new nodes to the trie. In existing nodes we just increase the `prefixCnt`. After the loop completes we need to increment the `wordCnt`. Note that we update the node elements `ch` and `str`. This IS NOT PART OF THE SOLUTION. It is only done to better understand what is going on.

        /**
         * Erases the string word from the trie.
         * It is guaranteed that for any function call to erase, 
         * the string word will exist in the trie.
         */
        public void erase(String word) {
            
            // ???? ????
            System.out.println("erase <<< word ==>" + word + "<==");

            // **** initialization ****
            Trie p          = this;
            boolean delete  = false;

            // **** traverse the word one character at a time ****
            for (char ch : word.toCharArray()) {

                // ???? ????
                System.out.println("erase <<<      p: " + p);

                // **** for ease of use ****
                int i = ch - 'a';

                // **** next child ****
                Trie next = p.children[i];  

                // ???? ????
                System.out.println("erase <<<   next: " + next);

                // **** decrement the prefix counter ****
                next.prefixCnt--;

                // **** flag if we need to delete the reference from children ****
                if (next.prefixCnt == 0) delete = true;

                // ????? ????
                System.out.println("erase <<< delete: " + delete);

                // **** delete the reference in this child ****
                if (delete) p.children[i] = null;

                // **** point to next node ****
                p = next;
            }

            // **** decrement the word counter ****
            p.wordCnt--;

            // ???? ????
            System.out.println("erase <<< p: " + p);
        }

This method is used to erase an existing (according to the requirements) word from the trie. It also 

The loop traverses the nodes in the trie for the characters in `word`.

We use the `next` variable to get to the next node in the trie which is associated with the value in `ch`. We decrement the `prefixCnt`. If the `prefixCnt` is equal to 0 we set the `delete` flag to true. This is used to delete the child node from the `children` Trie[] array in the `p` node. This releases the memory allocated for the node.

The last thing in the loop is to point to the next node.

After completing the loop, we end at the last node for the specified `word` so we need to decrement the `wordCnt`.

        /**
         * Returns the number of instances of the string word in the trie.
         */
        public int countWordsEqualTo(String word) {

            // **** initialization ****
            Trie p = this;

            // **** traverse the word one character at a time ****
            for (char ch : word.toCharArray()) {

                // **** for ease of use ****
                int i = ch - 'a';

                // **** word not found in trie ****
                if (p.children[i] == null) return 0;

                // **** move to next child ****
                p = p.children[i];
            }

            // **** count word in trie ****
            return p.wordCnt;
        }

This method is used to count the number of  `word` in the trie. It does so by traversing the trie one node at a time using the characters in `word`. The count is in the `wordCnt` member of the node in the last character of the specified word.

        /**
         * Returns the number of strings in the trie that have 
         * the string prefix as a prefix.
         */
        public int countWordsStartingWith(String prefix) {

            // **** initialization ****
            Trie p = this;

            // **** traverse prefix one character at a time ****
            for (char ch : prefix.toCharArray()) {

                // **** for ease of use ****
                int i = ch - 'a';

                // **** prefix not found ****
                if (p.children[i] == null) return 0;

                // **** move to next child ****
                p = p.children[i];
            }

            // **** ****
            return p.prefixCnt;
        }

This method counts the number of strings that contain the word in the `prefix` argument.

It does so by traversing the `prefix` one character at a time. If the node for the current character is not in the Trie[] array `children` the method returns 0 because the prefix was not found. Otherwise when the loop completes traversing the `prefix`, the result is in the `prefixCnt` of `p`.

As we have seen in the source code, the `ch` and `str` class members are only used to track the operation of the second implementation. Only the first three class members: `children`, `prefixCnt` and `wordCnt` are required.

Hope you enjoyed solving this problem as much as I did. The entire code for this project can be found in my GitHub repository named ImplementTrieIIPrefixTree.

Once again, sorry for the length of this post.

Please note that the code here presented might not be the best possible solution. After solving a problem, you should take a look at the best accepted solutions provided by the different websites (i.e., HackerRank, LeetCode to name a few). The idea is to learn concepts, not to memorize solutions.

If you have comments or questions regarding this, or any other post in this blog, please do not hesitate and leave me a note below. I will reply as soon as possible.

Keep on reading and experimenting. It is one of the best ways to learn, become proficient, refresh your knowledge and enhance your developer / engineering toolset.

Thanks for reading, feel free to connect with me John Canessa at LinkedIn.

Enjoy;

John

Leave a Reply

Your email address will not be published.

This site uses Akismet to reduce spam. Learn how your comment data is processed.