跳转至

667.键值映射 (Medium)*

题目描述*

标签*

字典树;

思路 & 代码*

sum 函数需要搜索前缀,那就是用字典树了。搜索到前缀后遍历那棵 N 叉树,计算键值总和。

class TrieNode {
public:
    int val;
    vector<TrieNode*> next;
    TrieNode() : val(0), next(vector<TrieNode*>(26, nullptr)) {}
    TrieNode(int _val) : val(_val), next(vector<TrieNode*>(26, nullptr)) {}
};
class MapSum {
private:
    TrieNode* root;
public:
    /** Initialize your data structure here. */
    MapSum() {
        root = new TrieNode();
    }

    void insert(string key, int val) {
        auto cur = root;
        for(auto&c : key) {
            if(cur->next[c - 'a'] == nullptr) {
                cur->next[c - 'a'] = new TrieNode();
            }
            cur = cur->next[c - 'a'];
        }
        cur->val = val;
    }

    int sum(string prefix) {
        auto cur = root;
        int res = 0;
        for(auto& c : prefix) {
            cur = cur->next[c - 'a'];
            if(cur == nullptr) {
                return 0;
            }
        }
        queue<TrieNode*> q;
        q.push(cur);
        while(!q.empty()) {
            auto node = q.front();
            q.pop();
            auto next = node->next;
            res += node->val;
            for(auto& p : next) {
                if(p != nullptr) {
                    q.push(p);
                }
            }
        }
        return res;
    }
};

/**
 * Your MapSum object will be instantiated and called as such:
 * MapSum* obj = new MapSum();
 * obj->insert(key,val);
 * int param_2 = obj->sum(prefix);
 */

最后更新: July 23, 2022