Implement Trie ll

Problem

Ninja has to implement a data structure ”TRIE” from scratch. Ninja has to complete some functions.

1) Trie(): Ninja has to initialize the object of this “TRIE” data structure.

2) insert(“WORD”): Ninja has to insert the string “WORD”  into this “TRIE” data structure.

3) countWordsEqualTo(“WORD”): Ninja has to return how many times this “WORD” is present in this “TRIE”.

4) countWordsStartingWith(“PREFIX”): Ninjas have to return how many words are there in this “TRIE” that have the string “PREFIX” as a prefix.

5) erase(“WORD”): Ninja has to delete one occurrence of the string “WORD” from the “TRIE”.

Note:

1. If erase(“WORD”) function is called then it is guaranteed that the “WORD” is present in the “TRIE”.

2. If you are going to use variables with dynamic memory allocation then you need to release the memory associated with them at the end of your solution.

Can you help Ninja implement the "TRIE" data structure?

Detailed explanation ( Input/output format, Notes, Images )

Constraints:

1 <= “T” <= 50
1 <= “N” <= 10000
 “WORD” = {a to z}
1 <= | “WORD” | <= 1000

Where “T” is the number of test cases, “N” denotes how many times the functions(as discussed above) we call, “WORD” denotes the string on which we have to perform all the operations as we discussed above, and | “WORD” | denotes the length of the string “WORD”.

Time limit: 1 sec.

Sample Input 1:

1
5
insert coding
insert ninja
countWordsEqualTo coding
countWordsStartingWith nin
erase coding

Sample Output 1:

1
1

Sample Input 2:

1
6
insert samsung
insert samsung
insert vivo
erase vivo
countWordsEqualTo samsung
countWordsStartingWith vi

Sample Output 2:

2
0

Solution

This is a continuation of the implementation of Trie I. Here, we also keep track of the count of words that start with a certain prefix and the total number of exact words. Instead of using a flag, the word count will track the number of words inserted into the trie.

    public class Node {
        Node[] nodeArr;
        int wordCount;
        int prefixCount;


        Node(){
            nodeArr = new Node[26];
            wordCount = 0;
            prefixCount = 0;
        }

        public void put(char chr, Node node){
            nodeArr[indexOf(chr)] = node;
        }

        public boolean containsKey(char chr){
            return nodeArr[indexOf(chr)] != null;
        }

        public Node getNext(char chr){
            return nodeArr[indexOf(chr)];
        }

        public void incrementWordCount(){
            this.wordCount += 1;
        }

        public void decrementWordCount(){
            this.wordCount -= 1;
        }

        public int getWordCount(){
            return this.wordCount;
        }

        public void incrementPrefixCount(){
            this.prefixCount += 1;
        }

        public void decrementPrefixCount(){
            this.prefixCount -= 1;
        }

        public int getPrefixCount(){
            return this.prefixCount;
        }

        private int indexOf(char chr){
            return chr-'a';
        }
    }

public class Trie {

    Node root;

    public Trie() {
        root = new Node();
    }

    public void insert(String word) {
        Node node = root;
        for(int i=0; i<word.length(); i++){
            char chr = word.charAt(i);
            if(!node.containsKey(chr)){
                node.put(chr, new Node());
            }
            node = node.getNext(chr);
            node.incrementPrefixCount();
        }
        node.incrementWordCount();
    }

    public int countWordsEqualTo(String word) {
        Node node = root;
        for(int i=0; i<word.length(); i++){
            char chr = word.charAt(i);
            if(!node.containsKey(chr)){
                return 0;
            }
            node = node.getNext(chr);
        }
        return node.getWordCount();
    }

    public int countWordsStartingWith(String word) {
        Node node = root;
        for(int i=0; i<word.length(); i++){
            char chr = word.charAt(i);
            if(!node.containsKey(chr)){
                return 0;
            }
            node = node.getNext(chr);
        }
        return node.getPrefixCount();
    }

    public void erase(String word) {
        Node node = root;
        for(int i=0; i<word.length(); i++){
            char chr = word.charAt(i);
            if(!node.containsKey(chr)){
                node.put(chr, new Node());
            }
            node = node.getNext(chr);
            node.decrementPrefixCount();
        }
        node.decrementWordCount();
    }
}