trie.hpp

Go to the documentation of this file.

#pragma once
 
#include <string>
#include <unordered_map>
#include <vector>
#include <algorithm>  // For std::transform


class TrieNode
{
public:
  std::unordered_map<char, TrieNode*> children;
  std::vector<int> indices;
  bool isEndOfWord;
 
  TrieNode() : isEndOfWord(false) {}
};

class Trie
{
private:
  TrieNode* root;
 
  // Helper function to convert characters to lowercase
  static auto toLowerCase(const std::string& str) -> std::string
  {
    std::string result = str;
    std::transform(result.begin(), result.end(), result.begin(), ::tolower);
    return result;
  }
 
  // Helper function to delete a word from the Trie recursively
  auto deleteWord(TrieNode* node, const std::string& word, int index, size_t depth = 0) -> bool
  {
    if (!node)
      return false;
 
    if (depth == word.size())
    {
      // Remove the index from this node
      auto& indices = node->indices;
      auto it = std::find(indices.begin(), indices.end(), index);
      if (it != indices.end())
      {
        indices.erase(it);
      }
      node->isEndOfWord = !indices.empty();
      return node->children.empty() && !node->isEndOfWord;
    }
 
    char c = toLowerCase(std::string(1, word[depth]))[0];
    if (deleteWord(node->children[c], word, index, depth + 1))
    {
      // Clean up the current node if it has no children and isn't the end of any word
      delete node->children[c];
      node->children.erase(c);
      return node->children.empty() && !node->isEndOfWord;
    }
    return false;
  }
 
public:
  Trie() : root(new TrieNode()) {}

  void insert(const std::string& word, int index)
  {
    TrieNode* node = root;
    std::string lower_word = toLowerCase(word);
 
    for (char c : lower_word)
    {
      if (!node->children.count(c))
      {
        node->children[c] = new TrieNode();
      }
      node = node->children[c];
      node->indices.push_back(index);
    }
 
    node->isEndOfWord = true;
  }

  auto search(const std::string& prefix) -> std::vector<int>
  {
    TrieNode* node = root;
    std::string lower_prefix = toLowerCase(prefix);
 
    for (char c : lower_prefix)
    {
      if (!node->children.count(c))
      {
        return {};  // No match found
      }
      node = node->children[c];
    }
 
    return node->indices;
  }

  void deleteWord(const std::string& word, int index)
  {
    deleteWord(root, word, index);
  }

  void clear()
  {
    delete root;
    root = new TrieNode();
  }

  [[nodiscard]] auto is_empty() const -> bool
  {
    return root->children.empty();
  }
};