Medium

题目描述

设计一个映射，允许你执行以下操作：

  • 将字符串键映射到给定值。
  • 返回具有等于给定字符串的前缀的键的值的总和。

实现 MapSum 类：

  • MapSum() 初始化 MapSum 对象。
  • void insert(String key, int val) 向映射中插入 key-val 键值对。如果键已经存在，原来的键值对将被新的键值对覆盖。
  • int sum(string prefix) 返回所有以 prefix 为前缀的键对应的值的总和。

示例 1：

输入：
["MapSum", "insert", "sum", "insert", "sum"]
[[], ["apple", 3], ["ap"], ["app", 2], ["ap"]]
输出：
[null, null, 3, null, 5]

解释：
MapSum mapSum = new MapSum();
mapSum.insert("apple", 3);  
mapSum.sum("ap");           // 返回 3 (apple = 3)
mapSum.insert("app", 2);    
mapSum.sum("ap");           // 返回 5 (apple + app = 3 + 2 = 5)

约束：

  • 1 <= key.length, prefix.length <= 50
  • key 和 prefix 仅由小写英文字母组成。
  • 1 <= val <= 1000
  • 最多调用 50 次 insert 和 sum。

解题思路

这道题有两种常见的解法：

方法一：哈希表 + 暴力搜索 最直接的思路是使用哈希表存储所有的键值对，在 sum 操作时遍历所有键，找出以给定前缀开头的键并累加它们的值。这种方法实现简单，但在前缀匹配时需要遍历所有键，时间复杂度较高。

方法二：字典树（Trie）（推荐） 更高效的解法是使用字典树。我们可以构建一个字典树，每个节点除了存储字符信息外，还存储从根节点到当前节点路径上所有完整单词的值的总和。这样在查询前缀和时，只需要走到前缀对应的节点，直接返回该节点存储的总和即可。

具体实现：

  • 每个 Trie 节点包含26个子节点（对应a-z）和一个 sum 值
  • insert 时，沿着路径更新每个节点的 sum 值
  • 如果是更新操作（key已存在），需要先减去旧值再加上新值
  • sum 时，沿着前缀路径找到对应节点，返回其 sum 值

字典树方法的优势在于 sum 操作的时间复杂度只依赖于前缀长度，而不是存储的键的数量。

代码实现

class MapSum {
private:
    struct TrieNode {
        TrieNode* children[26];
        int sum;
        
        TrieNode() : sum(0) {
            for (int i = 0; i < 26; i++) {
                children[i] = nullptr;
            }
        }
    };
    
    TrieNode* root;
    unordered_map<string, int> keyToVal;
    
public:
    MapSum() {
        root = new TrieNode();
    }
    
    void insert(string key, int val) {
        int oldVal = keyToVal.count(key) ? keyToVal[key] : 0;
        keyToVal[key] = val;
        int delta = val - oldVal;
        
        TrieNode* node = root;
        for (char c : key) {
            int idx = c - 'a';
            if (!node->children[idx]) {
                node->children[idx] = new TrieNode();
            }
            node = node->children[idx];
            node->sum += delta;
        }
    }
    
    int sum(string prefix) {
        TrieNode* node = root;
        for (char c : prefix) {
            int idx = c - 'a';
            if (!node->children[idx]) {
                return 0;
            }
            node = node->children[idx];
        }
        return node->sum;
    }
};

class MapSum:

    def __init__(self):
        self.trie = {}
        self.key_to_val = {}

    def insert(self, key: str, val: int) -> None:
        old_val = self.key_to_val.get(key, 0)
        self.key_to_val[key] = val
        delta = val - old_val
        
        node = self.trie
        for c in key:
            if c not in node:
                node[c] = {'sum': 0}
            node = node[c]
            node['sum'] += delta

    def sum(self, prefix: str) -> int:
        node = self.trie
        for c in prefix:
            if c not in node:
                return 0
            node = node[c]
        return node['sum']

public class MapSum {
    
    private class TrieNode {
        public TrieNode[] Children;
        public int Sum;
        
        public TrieNode() {
            Children = new TrieNode[26];
            Sum = 0;
        }
    }
    
    private TrieNode root;
    private Dictionary<string, int> keyToVal;

    public MapSum() {
        root = new TrieNode();
        keyToVal = new Dictionary<string, int>();
    }
    
    public void Insert(string key, int val) {
        int oldVal = keyToVal.ContainsKey(key) ? keyToVal[key] : 0;
        keyToVal[key] = val;
        int delta = val - oldVal;
        
        TrieNode node = root;
        foreach (char c in key) {
            int idx = c - 'a';
            if (node.Children[idx] == null) {
                node.Children[idx] = new TrieNode();
            }
            node = node.Children[idx];
            node.Sum += delta;
        }
    }
    
    public int Sum(string prefix) {
        TrieNode node = root;
        foreach (char c in prefix) {
            int idx = c - 'a';
            if (node.Children[idx] == null) {
                return 0;
            }
            node = node.Children[idx];
        }
        return node.Sum;
    }
}

var MapSum = function() {
    this.trie = {};
    this.keyToVal = new Map();
};

MapSum.prototype.insert = function(key, val) {
    const oldVal = this.keyToVal.get(key) || 0;
    this.keyToVal.set(key, val);
    const delta = val - oldVal;
    
    let node = this.trie;
    for (let c of key) {
        if (!(c in node)) {
            node[c] = {sum: 0};
        }
        node = node[c];
        node.sum += delta;
    }
};

MapSum.prototype.sum = function(prefix) {
    let node = this.trie;
    for (let c of prefix) {
        if (!(c in node)) {
            return 0;
        }
        node = node[c];
    }
    return node.sum;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

操作时间复杂度空间复杂度
insertO(m)O(m)
sumO(p)O(1)

其中：

  • m 是键的长度
  • p 是前缀的长度
  • 总空间复杂度：O(ALPHABET_SIZE × N × M)，其中 N 是插入的键的数量，M 是键的平均长度

相关题目