feat: implement additional functionality for Trie

Author: 4ndrelim

src/main/java/dataStructures/trie/Trie.java

@@ -1,12 +1,10 @@
 package dataStructures.trie;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
- * Implementation of Trie structure.
- * Supports the follwing common operations (see below for doc):
- * insert(String word)
- * search(String word)
- * startsWith(String prefix)
- * prune(String word)
+ * Implementation of a Trie; Here we consider strings (not case-sensitive)
  */
 public class Trie {
     private final TrieNode root;
@@ -15,99 +13,114 @@ public Trie() {
         root = new TrieNode();
     }
 
-    /**
-     * Insert a word into the trie; converts word to
-     * to lower-case characters before insertion.
-     *
-     * @param word the string to be inserted
-     */
     public void insert(String word) {
-        word = word.toLowerCase();
-        System.out.printf("~~~~~~~Inserting '%s'~~~~~~~%n", word);
-        TrieNode node = root;
+        word = word.toLowerCase(); // ignore case-sensitivity
+        TrieNode trav = root;
         for (int i = 0; i < word.length(); i++) {
             char curr = word.charAt(i);
-            if (!node.containsKey(curr)) {
-                node.insertKey(curr);
+            if (!trav.children.containsKey(curr)) {
+                trav.children.put(curr, new TrieNode()); // recall, the edges represent the characters
             }
-            node = node.getNext(curr); // go to the subsequent node!
+            trav = trav.children.get(curr);
         }
-        node.makeEnd();
+        trav.isEnd = true; // set word
     }
 
-    /**
-     * Search for a word (converted to lower-case) in the trie.
-     *
-     * @param word the string to look for
-     * @return boolean representing whether the word was found
-     */
     public boolean search(String word) {
-        word.toLowerCase();
-        System.out.printf("~~~~~~~Searching '%s'~~~~~~~%n", word);
-        TrieNode node = root;
+        word = word.toLowerCase();
+        TrieNode trav = root;
         for (int i = 0; i < word.length(); i++) {
             char curr = word.charAt(i);
-            if (node.containsKey(curr)) {
-                node = node.getNext(curr);
-            } else {
+            if (!trav.children.containsKey(curr)) {
                 return false;
             }
+            trav = trav.children.get(curr);
         }
-        return node.isEnd();
+        return trav.isEnd;
     }
 
-    /**
-     * Search for a prefix (converted to lower-case) in the trie.
-     * Note: very similar in implementation to search method
-     * except the search here does not need to look for end flag
-     *
-     * @param prefix the string to look for
-     * @return boolean representing whether the prefix exists
-     */
-    public boolean startsWith(String prefix) {
-        prefix = prefix.toLowerCase();
-        System.out.printf("~~~~~~~Looking for prefix '%s'~~~~~~~%n", prefix);
-        TrieNode node = root;
-        for (int i = 0; i < prefix.length(); i++) {
-            char curr = prefix.charAt(i);
-            if (node.containsKey(curr)) {
-                node = node.getNext(curr);
-            } else {
-                return false;
+    public void delete(String word) {
+        word = word.toLowerCase();
+        TrieNode trav = root;
+        for (int i = 0; i < word.length(); i++) {
+            char curr = word.charAt(i);
+            if (!trav.children.containsKey(curr)) {
+                return; // word does not exist in trie, so just return
             }
+            trav = trav.children.get(curr);
         }
-        return true;
+        trav.isEnd = false; // remove word from being tracked
     }
 
-    /**
-     * Removes a word from the trie by toggling the end flag;
-     * if any of the end nodes (next nodes relative to current)
-     * do not hold further characters, repetitively prune the trie
-     * by removing these nodes from the hashmap of the current node.
-     * Note: This method is useful in optimizing searching for a set of known words
-     * especially when the data to be traversed has words that are similar in spelling/
-     * repeated words which might have been previously found.
-     *
-     * @param word the word to be removed
-     */
-    public void prune(String word) {
-        word = word.toLowerCase();
-        System.out.printf("~~~~~~~Removing '%s'~~~~~~~%n", word);
-        TrieNode node = root;
-        TrieNode[] track = new TrieNode[word.length()];
+    // ABOVE ARE STANDARD METHODS OF A TYPICAL TRIE IMPLEMENTATION
+    // BELOW IMPLEMENTS TWO MORE COMMON / USEFUL METHODS FOR TRIE; IN PARTICULAR, NOTE THE PRUNING METHOD
+
+    public void deleteAndPrune(String word) {
+        List<TrieNode> trackNodes = new ArrayList<>();
+        TrieNode trav = root;
         for (int i = 0; i < word.length(); i++) {
             char curr = word.charAt(i);
-            track[i] = node;
-            node = node.getNext(curr);
+            if (!trav.children.containsKey(curr)) {
+                return; // word does not exist in trie
+            }
+            trackNodes.add(trav);
+            trav = trav.children.get(curr);
         }
-        node.removeEnd();
+        trav.isEnd = false;
+
+        // now we start pruning
         for (int i = word.length() - 1; i >= 0; i--) {
             char curr = word.charAt(i);
-            if (track[i].getNext(curr).getCharacters().size() > 0) {
-                break; // done further nodes are required
-            } else {
-                track[i].getCharacters().remove(curr);
+            TrieNode nodeBeforeCurr = trackNodes.get(i);
+            TrieNode nextNode = nodeBeforeCurr.children.get(curr);
+            if (nextNode.children.size() == 0) { // this node essentially doesn't track anything, remove
+                nodeBeforeCurr.children.remove(curr);
+            } else { // children.size() > 0; i.e. this node is still useful; no need to further prune upwards
+                break;
+            }
+        }
+        return;
+    }
+
+    public List<String> wordsWithPrefix(String prefix) {
+        List<String> ret = new ArrayList<>();
+        TrieNode trav = root;
+        for (int i = 0; i < prefix.length(); i++) {
+            char curr = prefix.charAt(i);
+            if (!trav.children.containsKey(curr)) {
+                return ret; // no words with this prefix
+            }
+            trav = trav.children.get(curr);
+        }
+        List<StringBuilder> allSuffix = getAllSuffixFromNode(trav);
+        for (StringBuilder sb : allSuffix) {
+            ret.add(prefix + sb.toString());
+        }
+        return ret;
+    }
+
+    public List<String> getAllWords() {
+        List<StringBuilder> allWords = getAllSuffixFromNode(root);
+        List<String> ret = new ArrayList<>();
+        for (StringBuilder sb : allWords) {
+            ret.add(sb.toString());
+        }
+        return ret;
+    }
+
+    private List<StringBuilder> getAllSuffixFromNode(TrieNode node) {
+        List<StringBuilder> ret = new ArrayList<>();
+        if (node.isEnd) {
+            ret.add(new StringBuilder(""));
+        }
+        for (char c : node.children.keySet()) {
+            TrieNode nextNode = node.children.get(c);
+            List<StringBuilder> allSuffix = getAllSuffixFromNode(nextNode);
+            for (StringBuilder sb : allSuffix) {
+                sb.insert(0, c); // insert c at the front
+                ret.add(sb);
             }
         }
+        return ret;
     }
 }

src/main/java/dataStructures/trie/TrieNode.java

@@ -1,85 +1,11 @@
 package dataStructures.trie;
 
-import java.util.HashMap;
 import java.util.Map;
 
 /**
- * Implementation of TrieNode for Trie structure (not restricited to alphabets)
- * Each TrieNode has the following attributes:
- * characters: HashMap whose keys represent the available characters and
- * values represent the nodes to select the next characters
- * NOTE: Presence of a character in the hashmap means that this character
- * is part of the trie structure and can be used to continue forming a searched word
- * or return the accumulation of letters as the word if the end flag is toggled to true
- * end       : Tells us whether the current TrieNode is actually a terminating flag
+ * Implementation of a TrieNode.
  */
-
 public class TrieNode {
-    private final Map<Character, TrieNode> characters;
-    private boolean end;
-
-    public TrieNode() {
-        characters = new HashMap<>();
-    }
-
-    /**
-     * Checks if node has a character.
-     *
-     * @param c character to check for presence
-     * @return boolean representing if the character exists
-     */
-    public boolean containsKey(char c) {
-        return characters.containsKey(c);
-    }
-
-    /**
-     * Get the next node at the index represented by the character.
-     *
-     * @param c character whose index holds the desired next node
-     * @return the desired node at that index
-     */
-    public TrieNode getNext(char c) {
-        return characters.get(c);
-    }
-
-    /**
-     * Get the avaialble characters can be used as the next letter from the current node.
-     *
-     * @return hashmap of characters and next nodes
-     */
-    public Map<Character, TrieNode> getCharacters() {
-        return characters;
-    }
-
-    /**
-     * Inserts a character to the current TrieNode.
-     *
-     * @param c character whose index represents where to insert the node
-     */
-    public void insertKey(char c) {
-        characters.put(c, new TrieNode());
-    }
-
-    /**
-     * Checks if the current TrieNode is a terminating flag.
-     *
-     * @return boolean value
-     */
-    public boolean isEnd() {
-        return end;
-    }
-
-    /**
-     * Make the current TrieNode a terminating flag/node.
-     */
-    public void makeEnd() {
-        end = true;
-    }
-
-    /**
-     * Remove terminating flag from current TrieNode.
-     */
-    public void removeEnd() {
-        end = false;
-    }
+    public Map<Character, TrieNode> children; // or an array of size 26 (assume not case-sensitive) to denote each char
+    public boolean isEnd; // the marker to indicate whether the path from the root to this node forms a known word
 }