Trie - Prefix Tree Implementation

doc.md

Trie (Prefix Tree) Problem Statement

Problem Title: Trie - Prefix Tree Implementation

Problem Description:

A Trie (pronounced "try") or prefix tree is a tree-like data structure that stores a dynamic set of strings. It is particularly useful for operations like searching, prefix matching, and autocomplete. Each node in the Trie represents a character of a string, and the path from the root to a node represents a prefix of one or more strings.

Your task is to implement a Trie class that supports the following operations:

1. Insert a String: Insert a word into the Trie.
2. Search for a String: Check if a word exists in the Trie.
3. Search for a Prefix: Check if any word in the Trie starts with a given prefix.
4. Delete a String: Remove a word from the Trie.
5. Count Words Starting with a Prefix: Return the number of words in the Trie that start with a given prefix.

Operations:

1. insert(word: str) -> None

   • Input: A string word to be inserted into the Trie.
   • Output: None.

2. search(word: str) -> bool

   • Input: A string word to search for in the Trie.
   • Output: True if the word exists in the Trie, otherwise False.

3. startsWith(prefix: str) -> bool

   • Input: A string prefix to check if any word in the Trie starts with it.
   • Output: True if there is any word in the Trie that starts with the prefix, otherwise False.

4. delete(word: str) -> bool

   • Input: A string word to delete from the Trie.
   • Output: True if the word was successfully deleted, otherwise False.

5. countWordsStartingWith(prefix: str) -> int

   • Input: A string prefix.
   • Output: The number of words in the Trie that start with the given prefix.

Constraints:

• The input words will consist of lowercase English letters.
• The length of each word will be at most 1000 characters.
• The number of operations (insert, search, startsWith, delete, countWordsStartingWith) will be at most 10^4.

Example Test Cases:

Test Case 1:

• Input:

  • Operations:
    1. insert("apple")
    2. search("apple")
    3. search("app")
    4. startsWith("app")
    5. insert("app")
    6. search("app")
    7. delete("apple")
    8. search("apple")
    9. search("app")

• Expected Output:

  • After operation 2: True
  • After operation 3: False
  • After operation 4: True
  • After operation 6: True
  • After operation 8: False
  • After operation 9: True

Test Case 2:

• Input:

  • Operations:
    1. insert("hello")
    2. insert("world")
    3. insert("hell")
    4. countWordsStartingWith("he")
    5. delete("hell")
    6. countWordsStartingWith("he")
    7. delete("hello")
    8. countWordsStartingWith("he")

• Expected Output:

  • After operation 4: 2
  • After operation 6: 1
  • After operation 8: 0

Test Case 3:

• Input:

  • Operations:
    1. insert("a")
    2. insert("aa")
    3. insert("aaa")
    4. countWordsStartingWith("a")
    5. countWordsStartingWith("aa")
    6. countWordsStartingWith("aaa")
    7. delete("aa")
    8. countWordsStartingWith("a")
    9. countWordsStartingWith("aa")

• Expected Output:

  • After operation 4: 3
  • After operation 5: 2
  • After operation 6: 1
  • After operation 8: 2
  • After operation 9: 1

Test Case 4:

• Input:

  • Operations:
    1. insert("cat")
    2. insert("car")
    3. insert("cart")
    4. insert("caterpillar")
    5. countWordsStartingWith("ca")
    6. delete("cart")
    7. countWordsStartingWith("ca")
    8. delete("caterpillar")
    9. countWordsStartingWith("ca")

• Expected Output:

  • After operation 5: 4
  • After operation 7: 3
  • After operation 9: 2

Test Case 5:

• Input:

  • Operations:
    1. insert("xyz")
    2. insert("xy")
    3. insert("x")
    4. countWordsStartingWith("x")
    5. delete("xy")
    6. countWordsStartingWith("x")
    7. delete("xyz")
    8. countWordsStartingWith("x")

• Expected Output:

  • After operation 4: 3
  • After operation 6: 2
  • After operation 8: 1

Challenge:

• Implement the Trie class with the above operations.
• Ensure that the operations are efficient, especially for large datasets.
• Handle edge cases, such as inserting an empty string, deleting a non-existent word, or counting words with an empty prefix.

Note: Do not provide the code, just the problem statement and test cases. The challenge is to implement the class based on the given operations and constraints.

---

class TrieNode:
    def __init__(self):
        # Each node holds a map from character → child node, and a flag for word‐end.
        self.children = {}
        self.is_end = False

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        """
        Insert `word` into the trie, iteratively.
        """
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True

    def search(self, word: str) -> bool:
        """
        Return True if `word` is in the trie (i.e. was inserted before).
        """
        node = self.root
        for ch in word:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return node.is_end

    def starts_with(self, prefix: str) -> bool:
        """
        Return True if there is any word in the trie that starts with `prefix`.
        """
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return True

    def delete(self, word: str) -> bool:
        """
        Delete `word` from the trie, if it exists. Returns True if deleted.
        This is done iteratively by first walking down and recording the path,
        then unmarking the end‐of‐word flag and pruning unused nodes bottom‐up.
        """
        node = self.root
        stack = []  # to remember (parent_node, char) on the path

        # 1) Walk down, storing the path
        for ch in word:
            if ch not in node.children:
                return False  # word not found
            stack.append((node, ch))
            node = node.children[ch]

        # 2) If this node wasn't marking end of a word, nothing to delete
        if not node.is_end:
            return False

        # 3) Unmark the word
        node.is_end = False

        # 4) Prune: walk the stored path in reverse
        #    and delete child nodes that are now useless
        for parent, ch in reversed(stack):
            child = parent.children[ch]
            # if child has no children and isn't end of another word → delete it
            if not child.children and not child.is_end:
                del parent.children[ch]
            else:
                break

        return True

# —— Example Usage ——
if __name__ == "__main__":
    trie = Trie()
    for word in ["apple", "app", "apt", "bat"]:
        trie.insert(word)

    print(trie.search("app"))       # True
    print(trie.search("apple"))     # True
    print(trie.search("ap"))        # False
    print(trie.starts_with("ap"))   # True
    print(trie.delete("app"))       # True
    print(trie.search("app"))       # False
    print(trie.search("apple"))     # True

---

Explanation of Operations

1. Node Structure

   • Each TrieNode has:
     • children: a dictionary mapping characters → child nodes.
     • is_end: a boolean flag indicating if a complete word ends at this node.

2. Insertion (insert)

   • Start at the root and for each character:
     • If the child for that character doesn’t exist, create it.
     • Move to that child.
   • After consuming all characters, mark the last node’s is_end = True.

3. Search (search)

   • Walk down from the root following each character in the query word.
   • If at any step the character isn’t found, return False.
   • After the loop, return the is_end status of the last node—only True if the exact word was inserted.

4. Prefix Check (starts_with)

   • Very similar to search, but you don’t require is_end = True at the end.
   • As long as you can follow the prefix all the way through, return True.

5. Deletion (delete)

   • Iterative walk: traverse the word while pushing (parent_node, ch) onto a stack.
   • If the word isn’t present or doesn’t terminate a word, you abort with False.
   • Unmark is_end on the terminal node.
   • Pruning: pop back up the stack, and at each step delete the child reference from its parent if:
     1. That child has no other children, and
     2. It isn’t marked as the end of another word.
   • Stop pruning once you hit a node that’s still needed (either has other children or is an is_end).

All methods avoid recursion by using simple loops and, for deletion, an explicit stack to backtrack. This keeps everything iterative while still maintaining full Trie functionality.

TrieTree.cc

/*
 *
 */

#include <iostream>
#include <unordered_map>
#include <vector>
#include <stack>
#include <queue>
#include <memory>

class TrieTree
{
private:
  struct TrieNode
  {
    std::unordered_map<char, std::unique_ptr<TrieNode>> child;
    bool is_end = false;
  };

  std::unique_ptr<TrieNode> (root) = nullptr;

public:
  TrieTree() {
    root = std::make_unique<TrieNode>();
  }

  void insert(const std::string& word) {
    TrieNode* curr = root.get();

    for (const char c : word)
    {
      auto it = curr->child.find(c);
      if (it != curr->child.end())
      {
        curr = it->second.get();
      }
      else
      {
        curr->child[c] = std::make_unique<TrieNode>();
        curr = curr->child[c].get();
      }
    }

    curr->is_end = true;
  }

  bool search(const std::string& word) const {
    TrieNode* curr = root.get();

    for (char c : word)
    {
      auto it = curr->child.find(c);
      if (it == curr->child.end())
        return false;

      curr = it->second.get();
    }

    return curr->is_end;
  }

  bool start_with(const std::string& word) const {
    TrieNode* curr = root.get();

    for (char c : word)
    {
      const auto it = curr->child.find(c);
      if (it == curr->child.end())
        return false;

      curr = it->second.get();
    }

    return true;
  }

  bool delete_node(const std::string& word) const {
    TrieNode* curr = root.get();
    std::stack<std::pair<char, TrieNode*>> st;

    for (const char& c : word)
    {
      auto it = curr->child.find(c);
      if (it == curr->child.end())
        return false;

      st.emplace(c, curr);
      curr = it->second.get();
    }

    // if found node is not word
    if (!curr->is_end)
      return false;

    curr->is_end = false;

    while (!st.empty())
    {
      auto [c, parent] = st.top();
      st.pop();

      auto child = parent->child[c].get();

      if (!child->is_end && child->child.empty())
      {
        parent->child.erase(c);
      }
      else
      {
        break;
      }
    }

    return true;
  }

  int count_words_starting_with(const std::string& word) const {
    TrieNode* curr = root.get();

    for (char c : word)
    {
      auto it = curr->child.find(c);
      if (it == curr->child.end())
        return 0;
      curr = curr->child[c].get();
    }

    std::queue<const TrieNode*> q;
    int count = 0;
    q.push(curr);

    while (!q.empty())
    {
      const TrieNode *parent = q.front();
      q.pop();

      if (parent->is_end)
        count++;

      for (const auto& [c, child] : parent->child)
          q.push(&*child);
    }

    return count;
  }
};

void test_1() {
  TrieTree tree;
  const std::vector<std::string> dict = {"apple", "app", "apt", "bat", "ab"};
  const std::vector<std::string> dict_del = {"apple", "ap", "app", "apt", "bat", "ball"};

  for (const std::string word : dict)
    tree.insert(word);

  std::cout << "Start With: " << tree.count_words_starting_with("a") << std::endl;

  std::cout << std::endl;

  for (const std::string word : dict_del)
    if (tree.search(word))
      std::cout << "Search: " << word << std::endl;
    else
      std::cout << "Not Search: " << word << std::endl;

  std::cout << std::endl;

  for (const std::string word : dict_del)
    if (tree.start_with(word))
      std::cout << "Prefix: " << word << std::endl;
    else
      std::cout << "Not Prefix: " << word << std::endl;

  std::cout << std::endl;

  for (const std::string word : dict_del)
    if (tree.delete_node(word))
      std::cout << "Delete: " << word << std::endl;
    else
      std::cout << "Not Delete: " << word << std::endl;

  std::cout << std::endl;

  for (const std::string word : dict_del)
    if (tree.search(word))
      std::cout << "Search: " << word << std::endl;
    else
      std::cout << "Not Search: " << word << std::endl;

  std::cout << std::endl;

  for (const std::string word : dict_del)
    if (tree.start_with(word))
      std::cout << "Prefix: " << word << std::endl;
    else
      std::cout << "Not Prefix: " << word << std::endl;
}

int main() {
  test_1();
  return 0;
}