Задача с собеседования в Amazon

Реализовать структуру данных MapSum, которая поддерживает две операции:
- insert(key, val) — вставить пару ключ-значение (или обновить значение, если ключ существует), key - строка, val - целое число
- sum(prefix) — вернуть сумму всех значений, ключи которых начинаются с заданного префикса

Пример:

mapSum.insert("apple", 3)
mapSum.sum("ap")         // return 3
mapSum.insert("app", 2)
mapSum.sum("ap")         // return 5 (apple + app = 3 + 2)

Работать должно линейно по времени

Наш чат алгоритмистов

Решение:

Будем использовать бор - префиксное дерево, где каждый узел хранит сумму всех значений с данным префиксом:
1. При вставке ключа вычисляем дельту (разницу между новым и старым значением)
2. Проходим по дереву вдоль символов ключа, создавая узлы при необходимости
3. В каждом узле увеличиваем значение на дельту
4. При запросе суммы просто идём по дереву до конца префикса и возвращаем значение узла

Дельта нужна для корректной обработки обновлений: если apple было 3, а стало 7, то во все узлы пути добавляем +4.

Это даёт O(k) по времени для обеих операций (где k — длина ключа/префикса) и O(n × k) по памяти (где n — количество ключей).

class TrieNode:
def init(self):
self.children = {}
self.value = 0

class MapSum:
def init(self):
self.root = TrieNode()
self.map = {} # Храним текущие значения для обработки обновлений

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

node = self.root
for char in key:
if char not in node.children:
node.children[char] = TrieNode()
node = node.children[char]
node.value += delta

def sum(self, prefix: str) -> int:
node = self.root
for char in prefix:
if char not in node.children:
return 0
node = node.children[char]
return node.value

Можно использовать просто словарь и при каждом sum() перебирать все ключи с помощью startswith(). Это даст O(1) для insert и O(n) для sum, что проще в реализации, но менее эффективно при частых запросах сумм.

@algoses