wildcard_trie.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <functional>
#include <memory>
#include <queue>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace SPSP
{
    template <typename TValue>
    class WildcardTrie
    {
        struct Node
        {
            TValue value;                                                   
            std::unordered_map<std::string, std::unique_ptr<Node>> childs;  
            size_t levelIndex = 0;                                          
            bool isLeaf = false;                                            
        };
 
        using BFSQueueT = std::queue<std::pair<std::string, const Node*>>;
 
        const std::string m_lSep;         
        const std::string m_lSingleWild;  
        const std::string m_lMultiWild;   
 
        Node m_root;                      
 
    public:
        WildcardTrie(const std::string& levelSeparator = "/",
                     const std::string& singleLevelWildcard = "+",
                     const std::string& multiLevelWildcard = "#") noexcept
            : m_lSep{levelSeparator}, m_lSingleWild{singleLevelWildcard},
              m_lMultiWild{multiLevelWildcard}
        {}
 
        TValue& operator[](const std::string& key)
        {
            Node* cur = &m_root;
            auto levels = this->splitToLevels(key);
 
            // Get or create child on each level
            for (size_t i = 0; i < levels.size(); i++) {
                auto& level = levels[i];
 
                // Create new child
                if (cur->childs.find(level) == cur->childs.end()) {
                    cur->childs[level] = std::make_unique<Node>();
                    cur->childs[level]->levelIndex = i + 1;
                }
 
                // Move to next level
                cur = cur->childs.at(level).get();
            }
 
            cur->isLeaf = true;
 
            return cur->value;
        }
 
        void insert(const std::string& key, const TValue& value)
        {
            (*this)[key] = value;
        }
 
        bool remove(const std::string& key)
        {
            Node* cur = &m_root;
            auto levels = this->splitToLevels(key);
 
            std::vector<Node*> nodeStack;
 
            // Get node if exists
            for (auto& level : levels) {
                nodeStack.push_back(cur);
 
                if (cur->childs.find(level) == cur->childs.end()) {
                    return false;
                }
                cur = cur->childs.at(level).get();
            }
 
            // Can't remove non-leaf node
            if (!cur->isLeaf) {
                return false;
            }
 
            cur->isLeaf = false;
 
            if (cur->childs.empty()) {
                // Delete all redundant ancestors
                // There is `int` instead of `size_t`, because we need signed type.
                for (int i = nodeStack.size() - 1; i >= 0; i--) {
                    Node* node = nodeStack.at(i);
                    if (node->isLeaf || node->childs.size() > 1 ||
                        node == &m_root) {
                        node->childs.erase(levels.at(i));
 
                        // Previous ancestors are no longer redundant 
                        break;
                    }
                }
            }
 
            return true;
        }
 
        using FindReturnT = std::unordered_map<std::string, const TValue&>;
 
        const FindReturnT find(const std::string& key) const
        {
            auto levels = this->splitToLevels(key);
 
            FindReturnT values;
 
            // Queue for to-be-processed nodes
            BFSQueueT nodeQueue;
            nodeQueue.push({ "", &m_root });
 
            while (!nodeQueue.empty()) {
                auto [nodeKey, node] = nodeQueue.front();
 
                if (node->levelIndex == levels.size() && node->isLeaf) {
                    // Match
                    values.insert({ nodeKey, node->value });
                }
                else if (node->levelIndex < levels.size()) {
                    // Enqueue relevant childs
                    for (auto& [childLevel, childNode] : node->childs) {
                        std::string childKey = nodeKey == ""
                            ? childLevel
                            : nodeKey + m_lSep + childLevel;
 
                        if (childLevel == levels.at(node->levelIndex) ||
                            childLevel == m_lSingleWild) {
                            // Key matches or has single-level wildcard
                            nodeQueue.push({ childKey, childNode.get() });
                        }
                        else if (childLevel == m_lMultiWild &&
                                 childNode->isLeaf) {
                            // Multi-level wildcard
                            values.insert({ childKey, childNode->value });
                        }
                    }
                }
 
                nodeQueue.pop();
            }
 
            return values;
        }
 
        void forEach(std::function<void(const std::string& key, const TValue& value)> f)
        {
            // Queue for to-be-processed nodes
            BFSQueueT nodeQueue;
            nodeQueue.push({ "", &m_root });
 
            while (!nodeQueue.empty()) {
                auto [nodeKey, node] = nodeQueue.front();
 
                // Call function
                if (node->isLeaf) {
                    f(nodeKey, node->value);
                }
 
                // Enqueue children
                for (auto& [childLevel, childNode] : node->childs) {
                    std::string childKey = nodeKey == ""
                        ? childLevel
                        : nodeKey + m_lSep + childLevel;
                    nodeQueue.push({ childKey, childNode.get() });
                }
 
                nodeQueue.pop();
            }
        }
 
        bool empty() const
        {
            return m_root.childs.empty();
        }
 
    protected:
        const std::vector<std::string> splitToLevels(const std::string& key) const
        {
            size_t curPos = 0, nextPos;
            std::vector<std::string> levels;
 
            while ((nextPos = key.find(m_lSep, curPos)) != std::string::npos) {
                levels.push_back(key.substr(curPos, nextPos - curPos));
                curPos = nextPos + m_lSep.length();
            }
 
            // Add the rest
            levels.push_back(key.substr(curPos));
 
            return levels;
        }
    };
} // namespace SPSP