Enhance Trie data structure with added methods and tests (#5538)

DIRECTORY.md

@@ -207,7 +207,7 @@
               * [SplayTree](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/SplayTree.java)
               * [Treap](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/Treap.java)
               * [TreeRandomNode](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/TreeRandomNode.java)
-              * [TrieImp](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/TrieImp.java)
+              * [Trie](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/Trie.java)
               * [VerticalOrderTraversal](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversal.java)
               * [ZigzagTraversal](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/trees/ZigzagTraversal.java)
           * devutils

src/main/java/com/thealgorithms/datastructures/trees/Trie.java

@@ -1,5 +1,28 @@
 package com.thealgorithms.datastructures.trees;
 
+import java.util.HashMap;
+
+/**
+ * Represents a Trie Node that stores a character and pointers to its children.
+ * Each node has a hashmap which can point to all possible characters.
+ * Each node also has a boolean value to indicate if it is the end of a word.
+ */
+class TrieNode {
+    char value;
+    HashMap<Character, TrieNode> child;
+    boolean end;
+
+    /**
+     * Constructor to initialize a TrieNode with an empty hashmap
+     * and set end to false.
+     */
+    TrieNode(char value) {
+        this.value = value;
+        this.child = new HashMap<>();
+        this.end = false;
+    }
+}
+
 /**
  * Trie Data structure implementation without any libraries.
  * <p>
@@ -14,27 +37,11 @@ package com.thealgorithms.datastructures.trees;
  * possible character.
  *
  * @author <a href="https://github.com/dheeraj92">Dheeraj Kumar Barnwal</a>
+ * @author <a href="https://github.com/sailok">Sailok Chinta</a>
  */
-public class TrieImp {
 
-    /**
-     * Represents a Trie Node that stores a character and pointers to its children.
-     * Each node has an array of 26 children (one for each letter from 'a' to 'z').
-     */
-    public class TrieNode {
-
-        TrieNode[] child;
-        boolean end;
-
-        /**
-         * Constructor to initialize a TrieNode with an empty child array
-         * and set end to false.
-         */
-        public TrieNode() {
-            child = new TrieNode[26];
-            end = false;
-        }
-    }
+public class Trie {
+    private static final char ROOT_NODE_VALUE = '*';
 
     private final TrieNode root;
 
@@ -42,8 +49,8 @@ public class TrieImp {
      * Constructor to initialize the Trie.
      * The root node is created but doesn't represent any character.
      */
-    public TrieImp() {
-        root = new TrieNode();
+    public Trie() {
+        root = new TrieNode(ROOT_NODE_VALUE);
     }
 
     /**
@@ -57,13 +64,15 @@ public class TrieImp {
     public void insert(String word) {
         TrieNode currentNode = root;
         for (int i = 0; i < word.length(); i++) {
-            TrieNode node = currentNode.child[word.charAt(i) - 'a'];
+            TrieNode node = currentNode.child.getOrDefault(word.charAt(i), null);
+
             if (node == null) {
-                node = new TrieNode();
-                currentNode.child[word.charAt(i) - 'a'] = node;
+                node = new TrieNode(word.charAt(i));
+                currentNode.child.put(word.charAt(i), node);
             }
             currentNode = node;
         }
+
         currentNode.end = true;
     }
 
@@ -80,13 +89,14 @@ public class TrieImp {
     public boolean search(String word) {
         TrieNode currentNode = root;
         for (int i = 0; i < word.length(); i++) {
-            char ch = word.charAt(i);
-            TrieNode node = currentNode.child[ch - 'a'];
+            TrieNode node = currentNode.child.getOrDefault(word.charAt(i), null);
+
             if (node == null) {
                 return false;
             }
             currentNode = node;
         }
+
         return currentNode.end;
     }
 
@@ -104,40 +114,89 @@ public class TrieImp {
     public boolean delete(String word) {
         TrieNode currentNode = root;
         for (int i = 0; i < word.length(); i++) {
-            char ch = word.charAt(i);
-            TrieNode node = currentNode.child[ch - 'a'];
+            TrieNode node = currentNode.child.getOrDefault(word.charAt(i), null);
             if (node == null) {
                 return false;
             }
+
             currentNode = node;
         }
+
         if (currentNode.end) {
             currentNode.end = false;
             return true;
         }
+
         return false;
     }
 
     /**
-     * Helper method to print a string to the console.
+     * Counts the number of words in the trie
+     *<p>
+     * The method traverses the Trie and counts the number of words.
      *
-     * @param print The string to be printed.
+     * @return count of words
      */
-    public static void sop(String print) {
-        System.out.println(print);
+    public int countWords() {
+        return countWords(root);
+    }
+
+    private int countWords(TrieNode node) {
+        if (node == null) {
+            return 0;
+        }
+
+        int count = 0;
+        if (node.end) {
+            count++;
+        }
+
+        for (TrieNode child : node.child.values()) {
+            count += countWords(child);
+        }
+
+        return count;
     }
 
     /**
-     * Validates if a given word contains only lowercase alphabetic characters
-     * (a-z).
-     * <p>
-     * The method uses a regular expression to check if the word matches the pattern
-     * of only lowercase letters.
+     * Check if the prefix exists in the trie
      *
-     * @param word The word to be validated.
-     * @return true if the word is valid (only a-z), false otherwise.
+     * @param prefix the prefix to be checked in the Trie
+     * @return true / false depending on the prefix if exists in the Trie
      */
-    public static boolean isValid(String word) {
-        return word.matches("^[a-z]+$");
+    public boolean startsWithPrefix(String prefix) {
+        TrieNode currentNode = root;
+
+        for (int i = 0; i < prefix.length(); i++) {
+            TrieNode node = currentNode.child.getOrDefault(prefix.charAt(i), null);
+            if (node == null) {
+                return false;
+            }
+
+            currentNode = node;
+        }
+
+        return true;
+    }
+
+    /**
+     * Count the number of words starting with the given prefix in the trie
+     *
+     * @param prefix the prefix to be checked in the Trie
+     * @return count of words
+     */
+    public int countWordsWithPrefix(String prefix) {
+        TrieNode currentNode = root;
+
+        for (int i = 0; i < prefix.length(); i++) {
+            TrieNode node = currentNode.child.getOrDefault(prefix.charAt(i), null);
+            if (node == null) {
+                return 0;
+            }
+
+            currentNode = node;
+        }
+
+        return countWords(currentNode);
     }
 }

src/test/java/com/thealgorithms/datastructures/trees/TrieTest.java

@@ -1,17 +1,21 @@
 package com.thealgorithms.datastructures.trees;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.List;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-public class TrieImpTest {
-    private TrieImp trie;
+public class TrieTest {
+    private static final List<String> WORDS = List.of("Apple", "App", "app", "APPLE");
+
+    private Trie trie;
 
     @BeforeEach
     public void setUp() {
-        trie = new TrieImp();
+        trie = new Trie();
     }
 
     @Test
@@ -66,11 +70,26 @@ public class TrieImpTest {
     }
 
     @Test
-    public void testIsValidWord() {
-        assertTrue(TrieImp.isValid("validword"), "Word should be valid (only lowercase letters).");
-        assertFalse(TrieImp.isValid("InvalidWord"), "Word should be invalid (contains uppercase letters).");
-        assertFalse(TrieImp.isValid("123abc"), "Word should be invalid (contains numbers).");
-        assertFalse(TrieImp.isValid("hello!"), "Word should be invalid (contains special characters).");
-        assertFalse(TrieImp.isValid(""), "Empty string should be invalid.");
+    public void testCountWords() {
+        Trie trie = createTrie();
+        assertEquals(WORDS.size(), trie.countWords(), "Count words should return the correct number of words.");
+    }
+
+    @Test
+    public void testStartsWithPrefix() {
+        Trie trie = createTrie();
+        assertTrue(trie.startsWithPrefix("App"), "Starts with prefix should return true.");
+    }
+
+    @Test
+    public void testCountWordsWithPrefix() {
+        Trie trie = createTrie();
+        assertEquals(2, trie.countWordsWithPrefix("App"), "Count words with prefix should return 2.");
+    }
+
+    private Trie createTrie() {
+        Trie trie = new Trie();
+        WORDS.forEach(trie::insert);
+        return trie;
     }
 }