main.cpp

#include <array>
#include <memory>
#include <functional>
#include <optional>

class TrieNode {
public:
    std::array<std::unique_ptr<TrieNode>, 26> next = {};
    bool has_word_end = false;
};

class Trie {
private:
    std::unique_ptr<TrieNode> root;
public:
    using NodeRef = std::reference_wrapper<const std::unique_ptr<TrieNode>>;
    Trie(): root(std::make_unique<TrieNode>()) {}

    // Inserts a word into the trie.
    void insert(string s) {
        auto p = std::ref(root);
        for(auto ch: s) {
            char idx = ch - 'a';
            if(p.get()->next[idx] == nullptr)
                p.get()->next[idx] = std::make_unique<TrieNode>();
            p = p.get()->next[idx];
        }
        p.get()->has_word_end = true;
    }

    // Returns if the word is in the trie.
    bool search(string key) {
        auto p = find(key);
        return p != std::nullopt && p->get()->has_word_end;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        return find(prefix) != std::nullopt;
    }

private:
    std::optional<NodeRef> find(string key) {
        auto p = std::cref(root);
        for(auto ch: key) {
            auto next = std::cref(p.get()->next[ch - 'a']);
            if (next.get() == nullptr) {
                return std::nullopt;
            }
            p = next;
        }
        return p.get();
    }
};