677. 键值映射-EXIAOHU Blog

题目链接：677. 键值映射

解法一

暴力遍历。插入的时间复杂度为$O(1)$；求和的时间复杂度为$O(N)$，其中$N$为键值对数目。

import typing


class MapSum:
    def __init__(self):
        self.data: typing.Dict[str, int] = dict()

    def insert(self, key: str, val: int) -> None:
        self.data[key] = val

    def sum(self, prefix: str) -> int:
        return sum(map(lambda item: item[1], filter(lambda item: item[0].startswith(prefix), self.data.items())))

解法二

字典树。

插入的时间复杂度为$O(len(key))$；求和的时间复杂度为$O(len(prefix))$。

由于本题中测试数据规模较小，字典树解法没有优势。

import typing

ALPHABET = 'abcdefghijklmnopqrstuvwxyz'


class DictTree:
    def __init__(self, val):
        self.next: typing.Dict[str, DictTree] = {}
        self.val = val

    def insert_and_get_child(self, char: str, val: int):
        if len(char) != 1 or char not in ALPHABET:
            raise ValueError()

        if self.next.get(char) is None:
            self.next[char] = DictTree(val)
        else:
            self.next[char].val += val

        return self.next[char]


class MapSum:
    def __init__(self):
        self.data = dict()
        self.tree = DictTree(0)

    def insert(self, key: str, val: int) -> None:
        old_val = self.data.get(key, 0)
        self.data[key] = val

        node = self.tree
        for char in key:
            node = node.insert_and_get_child(char, val-old_val)

    def sum(self, prefix: str) -> int:
        node = self.tree
        for char in prefix:
            node = node.insert_and_get_child(char, 0)
        return node.val