trie.cpp

#include <array>
#include <unordered_map>
#include <memory>
#include <iostream>
#include <string>

template<typename T>
using Ptr = std::shared_ptr<T>;

struct Node{
  char val;
  bool is_leaf;
  std::unordered_map<char, Ptr<Node>> children;

  Node(char c) : val(c), is_leaf(false) {};
};

Ptr<Node> create_node(char key){
  return std::make_shared<Node>(key);
}

Ptr<Node> get_child(Ptr<Node>& node, char ch){
  auto children = node->children;
  auto n = children.find(ch);

  if(n != children.end()){
    return n->second;
  }
  return nullptr;
}

void insert(Ptr<Node>& root, std::string word){
  Ptr<Node> node = root;
  for (char ch: word){
    Ptr<Node> child = get_child(node, ch);
    if(child != nullptr){
      node = child;
    }else{
      Ptr<Node> new_node = create_node(ch);
      node->children[ch] = new_node;
      node = new_node;
    }
  }
  node->is_leaf = true;
}

bool search(Ptr<Node>& root, std::string word){
  Ptr<Node> node = root;
  for (char ch: word){
    Ptr<Node> child = get_child(node, ch);
    if(child == nullptr){
      return false;
    }
    node = child;
  }
  return node->is_leaf;
}


void traverse(Ptr<Node>& root){
  if(root == nullptr){
    return;
  }

  std::cout << "\n";
  for(auto& child: root->children){
    auto node = child.second;
    std::cout << child.first << " ";
    if(node->is_leaf){
        std::cout << "\n";
    }
    traverse(node);
  }
}

void check_search(Ptr<Node>& root, std::string word){
  std::cout << word << ":" << search(root, word) << "\n";
}

int main(){
  auto root = create_node(' ');
  insert(root, "hello");
  insert(root, "eat");
  insert(root, "heat");

  traverse(root);

  check_search(root, "hello");
  check_search(root, "ello");
  check_search(root, "eate");
  check_search(root, "eat");
  check_search(root, "hheato");
  check_search(root, "heao");
  check_search(root, "heat");
}