0

Key-value store with prefix search

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: Key-value store with prefix search

Companies: Snowflake

Difficulty: Medium

Problem statement

Build an in-memory key-value store whose values can also be looked up by key prefix. Keys are strings and values are integers.

Implement the KVStore class:

  • KVStore() creates an empty store.
  • set(key, value) adds a new pair. Raise an error if key is already stored.
  • get(key) -> int returns the value for key, or -1 if key is not stored.
  • update(key, value) replaces the value of a stored key. Raise an error if key is not stored.
  • delete_key(key) removes key. Do nothing if it is not stored.
  • search(prefix) -> List[int] returns the values of all stored keys that start with prefix, ordered by key (lexicographic). An empty prefix matches every key. Return an empty list when nothing matches.

Prefix queries are frequent, so search should not scan the whole store. Its cost should depend on the prefix length and the size of the result, not on the total number of keys.

Constraints

  • Up to 10^5 operations in total.
  • Keys are 0 to 50 characters long; the empty string is a valid key.
  • Values fit in a 32-bit signed integer.

Example

store = KVStore()
store.set("pen", 4)
store.set("pencil", 2)
store.set("pet", 8)
store.set("map", 6)

store.search("pe")     -> [4, 2, 8]   # pen, pencil, pet
store.search("pen")    -> [4, 2]
store.get("pet")       -> 8
store.update("pen", 5)
store.delete_key("pencil")
store.search("pen")    -> [5]
store.get("pencil")    -> -1
store.search("z")      -> []
store.set("map", 1)    -> error: key exists

Interviewers often add parts one at a time, such as the sum of values under a prefix, returning only the first k matches, nested transactions with rollback, or making the store safe for concurrent use.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Clarify first. Pin these down before coding:

  • Errors. set on an existing key and update on a missing key raise. get on a missing key returns -1. delete_key on a missing key does nothing.
  • What "prefix search" returns. Here: the values of every key that starts with the prefix, ordered by key. Ask whether order matters, whether keys should come back too, and whether an empty prefix means "everything".
  • Read/write mix. If prefix queries are rare, a plain scan is fine. If they are frequent, or the result is small compared to the store, an index pays off.
  • Key alphabet and length. Lowercase ASCII or arbitrary Unicode? Short or long keys? This decides the trie node layout.

Brute force. Keep one dict. set, get, update, delete_key are O(1). search(p) scans every key, keeps the ones with key.startswith(p), and sorts them. That is O(n * |p| + m log m) per query, where m is the match count. It is correct and a good first answer, but each query touches the whole store even when it matches one key.

Optimal: hash map + trie. Keep the dict for point operations. Also insert every key into a trie, one character per edge, and store the value on the node where the key ends. To answer search(p):

  1. Walk |p| edges from the root. If an edge is missing, no key has that prefix.
  2. From the node you reach, do a depth-first walk that visits children in character order. Collect the value at every node that holds a key.

A pre-order walk with sorted children visits keys in lexicographic order, because a key comes before all its extensions, and siblings are ordered by their next character. So the result needs no extra sort.

Keeping the trie small. Each node stores count, the number of keys in its subtree. set adds 1 along the path; delete_key subtracts 1. When a child's count reaches 0, the whole branch below it is unused, so delete it and stop. This keeps memory proportional to the live keys, and it also gives count_prefix(p) in O(|p|) for free.

Invariants.

  • key in data if and only if the trie node at the end of key has has_value set, and both hold the same value.
  • Every node's count equals the number of stored keys in its subtree. The root's count equals len(data).
  • No trie node has count == 0, except the root.

Alternative: sorted key list. Keep keys in a sorted list and use bisect to find the range [p, p + ''). Queries are O(log n + m), but insert and delete are O(n) because of list shifting. A balanced tree or skip list (for example sortedcontainers.SortedList) makes all of them logarithmic. Mention this; it is what a database index does. The trie is the classic interview answer and is what the code uses.

Code

from __future__ import annotations


class _Node:
    __slots__ = ("children", "has_value", "value", "count")

    def __init__(self) -> None:
        self.children: dict[str, _Node] = {}
        self.has_value = False
        self.value = 0
        self.count = 0  # number of stored keys in this subtree


class KVStore:
    """Hash map for point lookups plus a trie for prefix queries."""

    def __init__(self) -> None:
        self._data: dict[str, int] = {}
        self._root = _Node()

    def set(self, key: str, value: int) -> None:
        if key in self._data:
            raise KeyError(f"key already exists: {key!r}")
        self._data[key] = value
        node = self._root
        node.count += 1
        for ch in key:
            node = node.children.setdefault(ch, _Node())
            node.count += 1
        node.has_value = True
        node.value = value

    def get(self, key: str) -> int:
        return self._data.get(key, -1)

    def update(self, key: str, value: int) -> None:
        if key not in self._data:
            raise KeyError(f"key not found: {key!r}")
        self._data[key] = value
        self._find(key).value = value  # type: ignore[union-attr]

    def delete_key(self, key: str) -> None:
        if key not in self._data:
            return  # deleting a missing key is a no-op
        del self._data[key]
        node = self._root
        node.count -= 1
        for ch in key:
            child = node.children[ch]
            child.count -= 1
            if child.count == 0:
                del node.children[ch]  # prune the now-empty branch
                return
            node = child
        node.has_value = False
        node.value = 0

    def search(self, prefix: str) -> list[int]:
        """Values of all keys starting with `prefix`, in key order."""
        start = self._find(prefix)
        if start is None:
            return []
        out: list[int] = []
        stack = [start]
        while stack:
            node = stack.pop()
            if node.has_value:
                out.append(node.value)
            # push children in reverse order so the smallest pops first
            for ch in sorted(node.children, reverse=True):
                stack.append(node.children[ch])
        return out

    def count_prefix(self, prefix: str) -> int:
        node = self._find(prefix)
        return node.count if node else 0

    def _find(self, prefix: str) -> _Node | None:
        node = self._root
        for ch in prefix:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

Complexity

Let L be the key length, |p| the prefix length, m the number of matching keys, S the number of trie nodes under the prefix node, and σ the alphabet size.

  • set: O(L) time (walk or create one node per character). Space O(L) new nodes in the worst case.
  • get: O(1) average, via the hash map.
  • update: O(L) time (hash map write plus walk to the trie node). If the trie node holds a reference to a shared value cell, this drops to O(1).
  • delete_key: O(L) time; it may free up to L nodes.
  • search: O(|p| + S log σ) time. S is at most the total length of the matching keys, so this is output-sensitive. Space O(m) for the result plus O(S) for the stack.
  • count_prefix: O(|p|).
  • Total space: O(n) for the map plus O(sum of key lengths) trie nodes.

Test cases

  • set("car",3), set("cart",7), set("cat",1), set("dog",9), set("ca",5):
    • search("ca") returns [5, 3, 7, 1] (keys ca, car, cart, cat).
    • search("car") returns [3, 7].
    • search("") returns [5, 3, 7, 1, 9], all values in key order.
    • search("x") returns [].
  • get("cat") returns 1; get("cow") returns -1.
  • update("cart", 8) then search("cart") returns [8].
  • set("car", 1) when car exists raises. update("cow", 1) when cow is missing raises.
  • delete_key("car"): get("car") returns -1, but search("car") still returns [8] because cart is kept.
  • delete_key("cart"): search("car") returns [], and count_prefix("ca") returns 2.
  • delete_key("nope") does nothing.
  • The empty key: set("", 4) works, get("") returns 4, and it comes first in search("").
  • Pruning: delete ca and cat so no c keys remain. The root has no c child left. Then set("car", 2) and search("c") returns [2].

Follow-ups

  1. Aggregate over a prefix (sum, max). Store subtree_sum on every node next to count. set, update and delete_key adjust it along the path by the change in value, so prefix_sum(p) is O(|p|). For max, a delete or a decreasing update cannot be undone by subtraction; recompute up the path from the children (O(L * σ)), or keep a heap per node.
  2. Top-k or pagination. Return only the first k matches, or the next page after a given key. The ordered DFS can stop after k results, so the cost is O(|p| + k * L). For "next page", walk to the cursor key and resume from there, or switch to a sorted structure with bisect.
  3. Transactions (begin / commit / rollback), possibly nested. Keep a stack of undo logs. Each write inside a transaction records the key's old state (absent, or old value). Rollback replays the log in reverse through the same set / update / delete_key paths, so the trie stays consistent. Commit merges the top log into the one below it. Reads need no changes because writes apply in place.
  4. Concurrency. A single lock around the store is correct and simple. For more read throughput, use a readers-writer lock, or shard the trie by first character so writers on different shards do not block each other. A copy-on-write trie (path copying) gives readers lock-free snapshots: a writer copies the O(L) nodes on the path and swaps in a new root.
  5. Memory and scale. Many long keys with shared prefixes waste nodes. Compress single-child chains into one edge labeled by a substring (a radix tree). For data larger than memory, keep keys sorted on disk (an LSM tree or B-tree) and answer prefix queries as range scans, which is how real key-value stores do it.

Interviewer expectations

Mid-level (L4)

  • Gets the four point operations right with the stated error behavior, using a hash map.
  • Recognizes the full scan as the brute force, then builds a working trie for search.
  • Tests the edge cases: missing keys, duplicate set, a prefix that is also a key.

Senior (L5)

  • Explains why a pre-order walk with sorted children returns keys in sorted order, so no final sort is needed.
  • Handles delete correctly, including pruning dead branches with subtree counts, and keeps the map and trie consistent.
  • Compares the trie with a sorted list or balanced tree, and picks based on the read/write mix and key shape.
  • Adds prefix aggregates or top-k with clean changes to the existing code.

Staff (L6+)

  • Drives the clarifying questions: ordering, result size, alphabet, read/write ratio, whether the store must be durable or concurrent.
  • Designs transactions or concurrency with clear invariants (undo logs, copy-on-write snapshots, sharding) and states their costs.
  • Connects the design to production systems: radix trees for memory, range scans over sorted storage for scale, and paging large prefix results.

Sign in to comment, vote or remix.