Simple Trie

Trie.h

#ifndef HD_TRIE_H
#define HD_TRIE_H

#include <map>
#include <string>
#include <vector>

/**
 * Trie data structure.
 *
 * Used for fast O(lg n) prefix searches.
 */
template <typename ValueT>
class Trie {
public:
    template <typename NodeValueT>
    class TrieNode {
    public:
        /**
         * Character representing this node.
         */
        wchar_t key;

        /**
         * Word value for this node (or nullptr if this an intermediate node).
         */
        std::shared_ptr<NodeValueT> value{};

        /**
         * Child nodes (by character).
         */
        std::map<const wchar_t, std::shared_ptr<TrieNode>> children;

        /**
         * @param key character representing this node.
         */
        explicit TrieNode(const wchar_t key) : key{key} {}
    };

    Trie() {
        root = std::make_shared<TrieNode<ValueT>>('\0');
    }

    ~Trie() {
        Clear();
    }

    /**
     * Clear the trie (at given node), resetting it to an empty state.
     *
     * @param node to clear from (or root if nullptr).
     */
    void Clear(std::shared_ptr<TrieNode<ValueT>> node = nullptr) {
        if (!node) {
            node = root;
        }

        for (auto &elem : node->children) {
            Clear(elem.second);
            elem.second = nullptr;
        }

        node->children.clear();
        node->value = nullptr;
    }

    /**
     * Get total number of values stored in the trie.
     *
     * @param node to start counting from.
     * @return value count.
     */
    [[nodiscard]] size_t GetCount(std::shared_ptr<TrieNode<ValueT>> node = nullptr) const {
        if (!node) {
            node = root;
        }

        size_t count{0};

        if (node->value) {
            count = 1;
        }

        for (const auto &elem : node->children) {
            count += GetCount(elem.second);
        }

        return count;
    }

    /**
     * Insert a value into the trie.
     *
     * @param term to insert.
     * @param value value to store in node.
     */
    void Insert(const std::string &term, std::shared_ptr<ValueT> const &value) {
        std::shared_ptr<TrieNode<ValueT>> currNode = root;

        for (auto ch : term) {
            auto childNode = currNode->children[ch];
            if (!childNode) {
                childNode = std::make_shared<TrieNode<ValueT>>(ch);
                currNode->children[ch] = childNode;
            }

            currNode = childNode;
        }

        currNode->value = value;
    }

    /**
     * Search the trie for a given term (prefix).
     *
     * @param term to search for.
     * @return search results.
     */
    [[nodiscard]] std::vector<std::shared_ptr<ValueT>> Search(const std::string &term) const {
        std::shared_ptr<TrieNode<ValueT>> currNode = root;

        for (auto ch : term) {
            if (!currNode->children.count(ch)) {
                // Empty vector
                return std::vector<std::shared_ptr<ValueT>>();
            }

            currNode = currNode->children.at(ch);
        }

        return Values(currNode);
    }

    /**
     * Get flat list of all words under the given node.
     *
     * @param node to get list for.
     * @return words.
     */
    [[nodiscard]] std::vector<std::shared_ptr<ValueT>> Values(std::shared_ptr<TrieNode<ValueT>> node = nullptr) const {
        if (!node) {
            node = root;
        }

        std::vector<std::shared_ptr<ValueT>> results;

        if (node->value) {
            results.push_back(node->value);
        }

        for (const auto &elem : node->children) {
            for (auto &childValue : Values(elem.second)) {
                results.push_back(childValue);
            }
        }

        return results;
    }

private:
    /**
     * Root node of trie.
     */
    std::shared_ptr<TrieNode<ValueT>> root;
};

#endif