题目

实现一个 MapSum 类里的两个方法,insertsum

对于方法 insert,你将得到一对(字符串,整数)的键值对。字符串表示键,整数表示值。如果键已经存在,那么原来的键值对将被替代成新的键值对。

对于方法 sum,你将得到一个表示前缀的字符串,你需要返回所有以该前缀开头的键的值的总和。

示例 1:

  1. 输入: insert("apple", 3), 输出: Null
  2. 输入: sum("ap"), 输出: 3
  3. 输入: insert("app", 2), 输出: Null
  4. 输入: sum("ap"), 输出: 5

方案一(空间换时间)

class MapSum:

    def __init__(self):
        """
        Initialize your data structure here.

        insert ("apple", 3) to self.lookup 
        {
            "a": {
                "sum": 3,
                "p": {
                    "sum": 3,
                    "p": {
                        "sum": 3,
                        "l": {
                            "sum": 3,
                            "e": {
                                "sum": 3,
                                "is_word": True
                            }
                        }
                    }
                }
            }
        }
        """
        self.lookup = {}
        self.keys = {} # "apple" -> 3

    def insert(self, key: str, val: int) -> None:
        node = self.lookup
        if key not in self.keys:
            self.keys[key] = val
            for c in key:
                if c not in node:
                    node[c] = {
                        "sum": val
                    }
                else:
                    node[c]["sum"] += val
                node = node[c]
        else:
            before_val = self.keys[key]
            diff = val - before_val
            for c in key:
                node[c]["sum"] += diff
                node = node[c]

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

        return node["sum"]

方案二

class MapSum:

    def __init__(self):
        """
        Initialize your data structure here.

        insert ("apple", 3) to self.lookup 
        {
            "a": {
                "p": {
                    "p": {
                        "l": {
                            "e": {
                                "sum": 3,
                            }
                        }
                    }
                }
            }
        }
        then insert ("app", 2) to self.lookup
        {
            "a": {
                "p": {
                    "p": {
                        "sum": 2,
                        "l": {
                            "e": {
                                "sum": 3,
                            }
                        }
                    }
                }
            }
        }
        """
        self.lookup = {}


    def insert(self, key: str, val: int) -> None:
        '''时间复杂度 O(n)'''
        node = self.lookup
        for c in key:
            if c not in node:
                node[c] = {}
            node = node[c]
        node["sum"] = val


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

        return self.count(node)


    def count(self, root) -> int:
        '''计算以 node 为根节点,下的 sum 的和'''
        _sum = 0
        for key in root:
            if key == "sum":
                _sum += root["sum"]
            else:
                _sum += self.count(root[key])

        return _sum

原文

https://leetcode-cn.com/explore/learn/card/trie/167/practical-application-i/647/