Problem Description

You need to design a data structure called MapSum that supports two operations:

1. Insert Operation: Store a key-value pair where the key is a string and the value is an integer. If the key already exists in the map, replace the old value with the new value.

2. Sum Operation: Given a prefix string, return the sum of all values whose keys start with that prefix.

The MapSum class should implement:

• MapSum(): Constructor that initializes an empty MapSum object
• insert(key, val): Adds or updates the key-value pair in the map
• sum(prefix): Returns the sum of all values whose keys begin with the given prefix

For example, if you insert ("apple", 3) and ("app", 2), then calling sum("ap") would return 5 because both "apple" and "app" start with "ap", and 3 + 2 = 5.

The solution uses a combination of a hash table and a Trie data structure. The hash table d keeps track of the current values for each key, while the Trie stores cumulative sums at each node representing prefixes. When inserting a key, the algorithm calculates the difference x between the new value and the existing value (if any), then updates all nodes along the path in the Trie by adding this difference. This ensures that each node maintains the correct sum for its prefix. The search operation traverses the Trie following the prefix path and returns the accumulated sum stored at the final node.

Intuition

The key insight is recognizing that we need to efficiently compute sums for all keys sharing a common prefix. A naive approach would be to iterate through all stored keys and check if each starts with the given prefix, but this would be inefficient.

This naturally leads us to think about Trie (prefix tree) data structure, which is specifically designed for prefix-based operations. In a standard Trie, each node represents a character in a path from root to leaf, and complete paths form words.

However, instead of just storing whether a word exists, we can augment each Trie node to store the sum of all values whose keys pass through that node. This way, when we reach the node representing a prefix, we immediately have the sum of all keys starting with that prefix.

The challenge arises when we need to handle updates - if a key already exists and we insert it again with a new value, we need to update not just one node but all nodes along the path from root to that key. Simply adding the new value would give incorrect results.

To solve this elegantly, we maintain a separate hash table d that tracks the current value for each key. When inserting, we calculate the difference x = newValue - oldValue (where oldValue is 0 if the key doesn't exist). We then traverse the Trie path and add this difference to each node's sum. This incremental update strategy ensures that:

• If it's a new key, we add the full value to each node on the path
• If it's an update, we only add the difference, effectively removing the old value and adding the new one in a single pass

This design gives us O(n) insertion time where n is the key length, and O(m) query time where m is the prefix length, making both operations very efficient.

Learn more about Trie patterns.

Solution Approach

The solution uses a combination of a hash table and a Trie to efficiently handle both insertions and prefix sum queries.

Data Structures

1. Hash Table (d): A dictionary that stores the current value for each key. This allows us to track what value is already associated with a key when we need to update it.

2. Trie: A prefix tree where each node contains:

   • children: An array of size 26 (for lowercase letters a-z) storing child nodes
   • val: The cumulative sum of all key-value pairs that pass through this node

Implementation Details

Trie Node Structure:

• Each node has a children array initialized with None values for all 26 letters
• The val field stores the sum of all values whose keys include this node in their path

Insert Operation in Trie:

def insert(self, w: str, x: int):
    node = self
    for c in w:
        idx = ord(c) - ord('a')  # Convert character to index (0-25)
        if node.children[idx] is None:
            node.children[idx] = Trie()  # Create new node if needed
        node = node.children[idx]
        node.val += x  # Add the difference to cumulative sum

The method traverses the Trie, creating nodes as needed, and adds the value x to each node along the path.

MapSum Insert Operation:

def insert(self, key: str, val: int) -> None:
    x = val - self.d[key]  # Calculate difference (0 if key doesn't exist)
    self.d[key] = val      # Update hash table with new value
    self.tree.insert(key, x)  # Insert difference into [Trie](/problems/trie_intro)

The clever part is calculating x = val - self.d[key]. Since defaultdict(int) returns 0 for non-existent keys:

• For new keys: x = val - 0 = val (adds the full value)
• For existing keys: x = newVal - oldVal (adds only the difference)

Search Operation:

def search(self, w: str) -> int:
    node = self
    for c in w:
        idx = ord(c) - ord('a')
        if node.children[idx] is None:
            return 0  # Prefix doesn't exist
        node = node.children[idx]
    return node.val  # Return cumulative sum at prefix node

The search traverses the Trie following the prefix path. If any character in the prefix doesn't exist, it returns 0. Otherwise, it returns the val stored at the final node, which represents the sum of all keys with this prefix.

Time Complexity

• Insert: O(n) where n is the length of the key
• Sum: O(m) where m is the length of the prefix

Space Complexity

• O(ALPHABET_SIZE × N × L) where ALPHABET_SIZE = 26, N is the number of unique keys, and L is the average length of keys

Example Walkthrough

Let's walk through a concrete example to illustrate how the MapSum data structure works:

Step 1: Initialize MapSum

mapSum = MapSum()
- Hash table d = {}
- Trie root initialized with empty children and val = 0

Step 2: Insert ("apple", 3)

- Calculate difference: x = 3 - 0 = 3 (key doesn't exist, so old value is 0)
- Update hash table: d = {"apple": 3}
- Insert into Trie with value 3:

  root
   |
   a (val=3)
   |
   p (val=3)
   |
   p (val=3)
   |
   l (val=3)
   |
   e (val=3)

Each node along the path gets val += 3.

Step 3: Insert ("app", 2)

- Calculate difference: x = 2 - 0 = 2 (new key)
- Update hash table: d = {"apple": 3, "app": 2}
- Insert into Trie with value 2:

  root
   |
   a (val=5)  // was 3, now 3+2=5
   |
   p (val=5)  // was 3, now 3+2=5
   |
   p (val=5)  // was 3, now 3+2=5
   |
   l (val=3)  // only "apple" passes through here
   |
   e (val=3)

Note: The first three nodes (a-p-p) are shared by both "apple" and "app", so their values are 5 (3+2).

Step 4: Query sum("ap")

- Traverse Trie: root → a → p
- Return value at node 'p': 5
- This correctly returns the sum of all keys starting with "ap" (apple=3 + app=2)

Step 5: Update by Insert ("apple", 5)

- Calculate difference: x = 5 - 3 = 2 (updating existing key from 3 to 5)
- Update hash table: d = {"apple": 5, "app": 2}
- Add difference 2 to Trie path for "apple":

  root
   |
   a (val=7)  // was 5, now 5+2=7
   |
   p (val=7)  // was 5, now 5+2=7
   |
   p (val=7)  // was 5, now 5+2=7
   |
   l (val=5)  // was 3, now 3+2=5
   |
   e (val=5)  // was 3, now 3+2=5

Step 6: Query sum("ap") again

- Traverse Trie: root → a → p
- Return value at node 'p': 7
- This correctly returns apple=5 + app=2 = 7

The key insight is that by storing cumulative sums at each node and using the difference calculation (new_value - old_value), we can efficiently handle both new insertions and updates while maintaining correct prefix sums throughout the Trie.

Solution Implementation

Time and Space Complexity

Time Complexity:

• insert(key, val): O(m) where m is the length of the key. The method traverses through each character of the key once to update the trie nodes.
• sum(prefix): O(p) where p is the length of the prefix. The method traverses through each character of the prefix to find the corresponding node.

Space Complexity:

• The overall space complexity is O(n × m × C), where:

  • n is the number of unique keys stored
  • m is the maximum length of the keys
  • C is the size of the character set (26 for lowercase English letters)

• The defaultdict stores n key-value pairs, contributing O(n × m) space (as keys are strings of length up to m)

• The Trie structure can have up to n × m nodes in the worst case (when all keys have completely different characters), and each node contains an array of 26 pointers to children, resulting in O(n × m × 26) space

• Therefore, the dominant factor is the Trie structure with O(n × m × C) where C = 26

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Handling of Key Updates

One of the most common mistakes is failing to properly handle the case when a key already exists and needs to be updated with a new value.

Pitfall Example:

def insert(self, key: str, val: int) -> None:
    # WRONG: Simply adding the new value without considering existing value
    self.trie.insert(key, val)  # This would add val on top of existing value
    self.key_value_map[key] = val

Why it's wrong: If "apple" was previously inserted with value 3, and we insert "apple" again with value 5, the trie nodes would accumulate both values (3 + 5 = 8) instead of replacing 3 with 5.

Correct Solution:

def insert(self, key: str, val: int) -> None:
    # Calculate the difference to add/subtract from trie nodes
    delta = val - self.key_value_map[key]  # defaultdict returns 0 for new keys
    self.key_value_map[key] = val
    self.trie.insert(key, delta)  # Only add the difference

2. Not Using defaultdict Leading to KeyError

Another common mistake is using a regular dictionary instead of defaultdict(int).

Pitfall Example:

def __init__(self):
    self.key_value_map = {}  # Regular dict instead of defaultdict
    self.trie = Trie()

def insert(self, key: str, val: int) -> None:
    # This will raise KeyError if key doesn't exist
    delta = val - self.key_value_map[key]  # KeyError for new keys!

Correct Solution:

def __init__(self):
    self.key_value_map = defaultdict(int)  # Returns 0 for missing keys
    # OR handle it explicitly:
    # self.key_value_map = {}

def insert(self, key: str, val: int) -> None:
    # If using regular dict, check existence first
    delta = val - self.key_value_map.get(key, 0)  # Use get() with default
    self.key_value_map[key] = val
    self.trie.insert(key, delta)

3. Accumulating Values at Wrong Trie Nodes

Some implementations mistakenly only update the value at the end node of a word, rather than updating all nodes along the path.

Pitfall Example:

def insert(self, word: str, delta: int) -> None:
    node = self
    for char in word:
        index = ord(char) - ord('a')
        if node.children[index] is None:
            node.children[index] = Trie()
        node = node.children[index]
    # WRONG: Only updating the final node
    node.val += delta  # Should update ALL nodes along the path

Why it's wrong: If we have "apple" with value 3 and "app" with value 2, searching for prefix "ap" would return 0 instead of 5, because the intermediate nodes weren't updated.

Correct Solution:

def insert(self, word: str, delta: int) -> None:
    node = self
    for char in word:
        index = ord(char) - ord('a')
        if node.children[index] is None:
            node.children[index] = Trie()
        node = node.children[index]
        node.val += delta  # Update EVERY node along the path

4. Memory Inefficiency with Fixed-Size Children Array

While not incorrect, using a fixed array of 26 elements for each node can be memory-inefficient when the trie is sparse.

Alternative Solution Using Dictionary:

class Trie:
    def __init__(self):
        self.children = {}  # Use dict instead of fixed array
        self.val = 0
  
    def insert(self, word: str, delta: int) -> None:
        node = self
        for char in word:
            if char not in node.children:
                node.children[char] = Trie()
            node = node.children[char]
            node.val += delta

This saves memory when most nodes don't have all 26 children, though it may have slightly worse cache locality.