Count of distinct substrings

Problem

Given a string s consisting of lowercase English characters, determine the total number of distinct non-empty substrings present in the string. A substring is defined as a contiguous block of characters within the string.

Two substrings are considered distinct if their contents differ, even if they originate from different positions in the string. (link)

Note: The empty substring is not counted.

Examples :

Input: s = "ababa"
Output: 9
Explanation: All distinct substrings of "ababa" are: "a", "b", "ab", "ba", "aba", "bab", "abab", "baba", "ababa".

Input: s = "aaa"
Output: 3
Explanation: The distinct substrings of "aaa" are: "a", "aa", "aaa".

Constraints:
1 ≤ s.size() ≤ 3000

Solution

We will use tries to find distinct substrings. Each time we insert a node in the trie, the substring count increases.

Time - O(nxn)

class Node {
    Node[] nodeArr;

    Node(){
        this.nodeArr = new Node[26];
    }

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

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

    public Node get(char chr){
        return this.nodeArr[indexOf(chr)];
    }

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

}

class Trie {

    Node root;
    int distinctWords;

    Trie() {
        this.root = new Node();
        this.distinctWords = 0;
    }

    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());
                this.distinctWords+=1;
            }
            node = node.get(chr);
        }
    }
}

class Solution {

    public static int countSubs(String s) {

        Trie trie = new Trie();
        for(int i=0; i<s.length(); i++){
            String word = "";
            for(int j=i; j<s.length(); j++){
                word += s.charAt(j);
                trie.insert(word);
            }
        }

        return trie.distinctWords;
    }
}