Add Another Trie Test (WordDictionary)

src/main/java/org/example/Trie.java

@@ -20,37 +20,31 @@ public void insert(String word) {
         node.isEndOfWord = true;
     }
 
+
     public boolean search(String word) {
-        TrieNode node = root;
-        for (int i = 0; i < word.length(); i++) {
-            if (node == null) {
-                return false;
-            }
-            if (!node.contains(word.charAt(i))) {
-                return false;
-            }
-            node = node.getSubNode(word.charAt(i));
-            if (i == word.length() - 1 && node.isEndOfWord) {
-                return true;
-            }
-        }
-        return false;
+        TrieNode node = searchPrefix(word);
+        return node != null && node.isEndOfWord;
     }
 
     public boolean startsWith(String word) {
+        return searchPrefix(word) != null;
+    }
+
+    public TrieNode searchPrefix(String word) {
         TrieNode node = root;
         for (int i = 0; i < word.length(); i++) {
             if (node == null) {
-                return false;
+                return null;
             }
             if (!node.contains(word.charAt(i))) {
-                return false;
+                return null;
             }
             node = node.getSubNode(word.charAt(i));
         }
-        return true;
+        return node;
     }
 
+
     @Override
     public String toString() {
         return "Trie{" +

src/main/java/org/example/TrieNode.java

@@ -1,6 +1,7 @@
 package org.example;
 
 import java.util.Arrays;
+import java.util.Objects;
 
 public class TrieNode {
     private static final int MAX_SIZE = 26;
@@ -24,6 +25,10 @@ public void put(char ch, TrieNode node) {
         subNodes[ch - 'a'] = node;
     }
 
+    public TrieNode[] getExistingSubNodes() {
+        return Arrays.stream(subNodes).filter(Objects::nonNull).toArray(TrieNode[]::new);
+    }
+
     @Override
     public String toString() {
         return "TrieNode{sub=" + Arrays.toString(subNodes) + "}";

src/main/java/org/example/WordDictionary.java

@@ -0,0 +1,53 @@
+package org.example;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class WordDictionary {
+    private Node root;
+
+    public WordDictionary() {
+        root = new Node();
+    }
+
+    public void addWord(String word) {
+        Node node = root;
+        for (char ch : word.toCharArray()) {
+            if (!node.subNodes.containsKey(ch)) {
+                node.subNodes.put(ch, new Node());
+            }
+            node = node.subNodes.get(ch);
+        }
+        node.isEndOfWord = true;
+    }
+
+    public boolean search(String word) {
+        return helper(root, word);
+    }
+
+    public boolean helper(Node node, String word) {
+        for (int i = 0; i < word.length(); i++) {
+            char ch = word.charAt(i);
+            if (!node.subNodes.containsKey(ch)) {
+                if (ch == '.') {
+                    for (char x : node.subNodes.keySet()) {
+                        Node child = node.subNodes.get(x);
+                        if (helper(child, word.substring(i + 1))) {
+                            return true;
+                        }
+                    }
+                }
+                return false;
+            } else {
+                node = node.subNodes.get(ch);
+            }
+        }
+        return node.isEndOfWord;
+    }
+
+    private class Node {
+        Map<Character, Node> subNodes = new HashMap<>();
+        boolean isEndOfWord;
+    }
+
+}

src/test/java/WordSearchTest.java

@@ -0,0 +1,59 @@
+import org.example.Trie;
+import org.example.WordDictionary;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class WordSearchTest {
+
+    // https://leetcode.com/problems/implement-trie-prefix-tree/
+    @Test
+    void trie() {
+//        TrieNode node = TrieNodeSupport.linkLeftsInOrder("apple");
+//        System.out.println("node = " + node);
+        Trie trie = new Trie();
+        trie.insert("apple");
+        trie.insert("app");
+        boolean search = trie.search("app");
+        System.out.println("search = " + search);
+        boolean startsWith = trie.startsWith("apple");
+        System.out.println("startsWith = " + startsWith);
+    }
+
+    // https://leetcode.com/problems/design-add-and-search-words-data-structure/
+    @Test
+    void wordDictionary() {
+// "addWord", ["at"],
+// "addWord", ["and"],
+// "addWord", ["an"],
+// "addWord", ["add"],
+// "search", ["a"],
+// "search", [".at"],
+// "addWord", ["bat"],
+// "search", [".at"],
+// "search", ["an."],
+// "search", ["a.d."],
+// "search", ["b."],
+// "search", ["a.d"],
+// "search"] ["."]]
+
+        WordDictionary wordDictionary = new WordDictionary();
+        wordDictionary.addWord("at");
+        wordDictionary.addWord("and");
+        wordDictionary.addWord("an");
+        wordDictionary.addWord("add");
+        List<Boolean> result = new ArrayList<>();
+        result.add(wordDictionary.search("a")); // false
+        result.add(wordDictionary.search(".at")); // false
+        wordDictionary.addWord("bat");
+        result.add(wordDictionary.search(".at")); // true
+        result.add(wordDictionary.search("an.")); // true
+        result.add(wordDictionary.search("a.d.")); // false
+        result.add(wordDictionary.search("b.")); // false
+        result.add(wordDictionary.search("a.d")); // true
+        result.add(wordDictionary.search(".")); // false
+
+        result.forEach(System.out::println);
+    }
+}