Map Sum Pairs (LeetCode 677)

The problem

Design a MapSum class:

• insert(key, val) stores key → val. If the key already exists, the new value replaces the old one.
• sum(prefix) returns the total of the values of every key that starts with prefix.

insert("apple", 3)
sum("ap")          -> 3
insert("app", 2)
sum("ap")          -> 5

After the second insert, both apple (3) and app (2) start with ap, so the sum is 5.

How to think about it

Start with brute force. Keep a plain dict of keys to values. For sum(prefix), loop over every key and add the value if key.startswith(prefix). That is O(N·L) per sum for N keys. It is simple and correct, and it already handles overwrites, because a dict keeps one value per key.

The key insight: a trie finds the prefix, and a value map handles overwrites. The trie answers “does any key start with this prefix?” in O(L): if the walk falls off, the answer is 0 straight away. The separate vals dict keeps exactly one value per key, so re-inserting a key never double counts. The solution below then totals the matching keys from vals, which keeps the code short and obviously correct.

The fully optimised version. To make sum O(L) as well, store a running total at every trie node. On insert, compute delta = val - old_val and add delta to each node on the key’s path. Then sum(prefix) is just the total stored at the prefix’s node. Mention this in an interview; it is the reason a trie fits this problem.

Which pattern is this? Trie with data on the nodes. The tell is “aggregate over all keys sharing a prefix”.

Approach

1. Keep vals = {} (key to latest value) and a trie root = {}.
2. insert(key, val): set vals[key] = val (this overwrites cleanly), then walk the key into the trie, creating nodes.
3. sum(prefix): walk the prefix down the trie. If a character is missing, return 0.
4. Otherwise, add up vals[k] for every key k that starts with prefix.

Solution

class MapSum:
    def __init__(self): self.vals = {}; self.root = {}
    def insert(self, key, val):
        self.vals[key] = val; n = self.root            # overwrite, never add twice
        for c in key: n = n.setdefault(c, {})
    def sum(self, prefix):
        n = self.root
        for c in prefix:
            if c not in n: return 0                    # no key has this prefix
            n = n[c]
        return sum(v for k, v in self.vals.items() if k.startswith(prefix))
# insert("apple",3); sum("ap") -> 3; insert("app",2); sum("ap") -> 5

Dry run

Running the example, plus an overwrite and two extra queries:

The overwrite changes apple from 3 to 5, so the sum becomes 7, not 10.

Complexity

• Time: insert is O(L). One dict write plus one walk of the key. sum is O(L + N·L) as written: the walk, then a scan of all N keys. With per-node totals, sum drops to O(L).
• Space: O(S) for the trie and the value map, where S is the total length of all keys.

Edge cases and pitfalls

• Overwrites. Adding val to every node on insert, without subtracting the old value, double counts a re-inserted key. Use the delta.
• Prefix with no keys returns 0 from the walk.
• Prefix equal to a whole key (sum("apple")) includes that key itself.
• Empty prefix sums every value.

Related problems

• Implement Trie (Prefix Tree) (LeetCode 208): the plain trie this extends.
• Search Suggestions System (LeetCode 1268): another “everything under this prefix” query.
• Design HashMap (LeetCode 706): the key-value side of this design.