0

Shard rebalancing with a per-key replica limit

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: Shard rebalancing with a per-key replica limit

Companies: OpenAI

Difficulty: Medium

Problem statement

Our key-value store spreads keys across shards. Each shard is given as a string "<id>:<start>:<end>" and owns every integer key in the inclusive range [start, end]. Shards may overlap. The depth of a key is the number of shards that cover it.

Implement:

def rebalance(limit: int, shards: list[str]) -> list[str]: ...

Return the rebalanced shards in the same "id:start:end" format, sorted by (start, id), so that:

  1. No key has depth above limit (limit >= 1).
  2. Earlier shards win. Process shards in order of (start, id). A shard placed earlier is never changed by a later one. A later shard resolves a conflict only by moving its start to the right, as little as possible. Its end stays. A shard left empty (start > end) is dropped.
  3. No gaps. In the result, every key from the smallest original start to the largest original end is covered. Close a hole by extending the end of the kept shard that reaches furthest to its left, up to the key before the next shard.

You may not create new shard IDs or split a shard into two ranges. Keys fit in 64-bit integers and can be as large as 10^9. There are up to a few thousand shards.

Example 1

limit = 2
shards = ["x:0:50", "y:10:60", "z:20:70"]

Keys 20..50 are covered by three shards. x and y come first and stay. z moves its start to 51.

["x:0:50", "y:10:60", "z:51:70"]

Example 2

limit = 1
shards = ["a:0:10", "b:5:20", "c:30:40", "d:35:50"]

b moves to 11:20 and d moves to 41:50. Nothing covers keys 21..29, so b is extended to 29.

["a:0:10", "b:11:29", "c:30:40", "d:41:50"]

Approach

Each shard owns an inclusive range of integer keys. The depth of a key is the number of shards that cover it. After rebalance(limit, shards):

  • no key has depth above limit;
  • every key from the smallest start to the largest end is covered (depth at least 1);
  • shard IDs are kept, and no shard is split into two pieces.

The rules that make the answer unique:

  1. Process shards in order of (start, id). An earlier shard has priority and is never changed by a later one.
  2. A later shard resolves a conflict only by moving its start to the right. Its end stays. If it becomes empty (start > end), drop it.
  3. After trimming, close each hole between covered ranges by extending the end of the shard that reaches furthest left of the hole.

Brute force. Fill a dictionary depth[key] and walk each shard key by key, skipping keys that are full. This is O(n * R), where R is the key span. The keys are arbitrary integers, so R can be 10^9. Say why that fails, then switch to working with interval endpoints.

Greedy on endpoints. Take the shards in priority order. Keep the shards already placed in kept. For the current shard [s, e]:

  • Clip every kept shard to [s, e] and turn each clip into two events: +1 at its left edge and -1 one past its right edge.
  • Sort the events and sweep them, tracking depth. Remember the rightmost key in [s, e] whose depth is already limit.
  • If there is such a key p, set s = p + 1. Keys at or left of p cannot take this shard, and the shard must stay one contiguous range, so it has to start after p. Moving the start to p + 1 and no further keeps as much of the shard as possible.

Why the result is valid: every key in the new [p + 1, e] had depth at most limit - 1 before this shard was added, so adding it brings depth to at most limit. Kept shards never change again, so the bound holds to the end.

Why coverage is kept: we only drop keys that are already covered limit >= 1 times by earlier shards. Trimming never opens a new hole. The only holes are gaps that were already in the input.

Filling gaps. Sort the kept shards by start. Track reach, the kept shard with the largest end so far. If the next shard starts after reach.end + 1, nothing covers the keys in between, so extending reach.end to next.start - 1 brings their depth from 0 to 1. Because limit >= 1, this is always allowed.

Invariants.

  • The kept list always satisfies depth <= limit at every key.
  • A shard's start never decreases. Its end only grows, and only into keys that had depth 0.
  • The output is sorted by (start, id) and uses the "id:start:end" format.

Code

from dataclasses import dataclass
from typing import List, Optional


@dataclass
class Shard:
    id: str
    start: int
    end: int

    @staticmethod
    def parse(text: str) -> "Shard":
        sid, start, end = text.rsplit(":", 2)
        return Shard(sid, int(start), int(end))

    def __str__(self) -> str:
        return f"{self.id}:{self.start}:{self.end}"


def rightmost_saturated(lo: int, hi: int, kept: List[Shard], limit: int) -> Optional[int]:
    """Largest key in [lo, hi] already covered by >= limit kept shards."""
    events: List[tuple] = []
    for k in kept:
        left, right = max(lo, k.start), min(hi, k.end)
        if left <= right:
            events.append((left, 1))
            events.append((right + 1, -1))
    events.sort()  # at equal keys, -1 sorts before +1
    best, depth, i = None, 0, 0
    while i < len(events):
        key = events[i][0]
        while i < len(events) and events[i][0] == key:
            depth += events[i][1]
            i += 1
        if depth >= limit and i < len(events):
            # depth holds on [key, next_key - 1]; a run always ends at a -1 event
            best = events[i][0] - 1
    return best


def rebalance(limit: int, shards: List[str]) -> List[str]:
    if limit < 1:
        raise ValueError("limit must be >= 1")
    items = sorted((Shard.parse(s) for s in shards), key=lambda s: (s.start, s.id))

    # Phase 1: trim. Earlier shards win; a later shard only moves its start right.
    kept: List[Shard] = []
    for sh in items:
        sat = rightmost_saturated(sh.start, sh.end, kept, limit)
        if sat is not None:
            sh.start = sat + 1
        if sh.start <= sh.end:
            kept.append(sh)

    # Phase 2: fill gaps so [min start, max end] is covered with no holes.
    kept.sort(key=lambda s: (s.start, s.id))
    reach: Optional[Shard] = None  # kept shard with the largest end so far
    for sh in kept:
        if reach is not None and sh.start > reach.end + 1:
            reach.end = sh.start - 1  # the gap had depth 0, so depth 1 is legal
        if reach is None or sh.end > reach.end:
            reach = sh
    return [str(s) for s in kept]

Complexity

With n shards:

  • Parsing and sorting: O(n log n).
  • Trimming: each shard sweeps up to 2n events from the shards already kept, sorted in O(n log n). That is O(n^2 log n) in total. It does not depend on the key span, so [0, 10^9] costs the same as [0, 10].
  • Gap filling: one pass after a sort, O(n log n).
  • Space: O(n) for the kept list and one event list.

The quadratic step is easy to explain and fine for the few thousand shards a real cluster has. Follow-up 1 makes it O(n log n).

Test cases

  • limit=1, ["a:0:100", "b:80:180"] -> ["a:0:100", "b:101:180"] (the overlap 80..100 leaves b).
  • limit=2, ["x:0:50", "y:10:60", "z:20:70"] -> ["x:0:50", "y:10:60", "z:51:70"] (keys 20..50 had depth 3).
  • limit=1, ["a:0:10", "b:3:5"] -> ["a:0:10"] (b lies inside a and is dropped).
  • limit=1, ["a:0:10", "b:20:30"] -> ["a:0:19", "b:20:30"] (the gap 11..19 goes to a).
  • limit=3, ["a:0:10", "b:0:10"] -> unchanged (depth 2 is under the limit).
  • limit=1, ["b:5:9", "a:5:9"] -> ["a:5:9"] (equal starts: the ID breaks the tie).
  • limit=1, ["a:0:10", "b:5:20", "c:30:40", "d:35:50"] -> ["a:0:10", "b:11:29", "c:30:40", "d:41:50"] (trim, then fill the gap 21..29 by extending b).
  • limit=2, ["a:0:3", "c:0:8", "b:2:8", "d:3:9"] -> ["a:0:3", "c:0:8", "b:4:8", "d:9:9"] (d is pushed past everything that is full and keeps one key).
  • limit=1, [] -> [].
  • limit=0 -> ValueError.
  • A randomized check against a key-by-key brute force: every key has depth <= limit, the union equals [min start, max end], and no start moved left.

Follow-ups

  1. Make it O(n log n). Compress all start and end+1 coordinates. Build a segment tree over the compressed segments that supports range add and "rightmost index in [l, r] whose value is >= limit" (store the max per node and descend right first). Each shard does one query and one range add, each O(log n).
  2. Online add_shard / remove_shard. Keep shards in a sorted container keyed by (original start, id) and remember each shard's original range. A new shard placed late in the order needs one query. One placed early can push all later shards, so re-run the greedy from that position. Removing a shard can let later shards grow back to their original starts, so recompute from its position as well. Point out that "earlier wins" makes changes spread to the right.
  3. Minimize data movement. Moving a start means migrating those keys. Report the moved ranges ([old_start, new_start - 1] per shard, plus gap extensions) as a migration plan. Then ask whether to optimize "fewest keys moved" instead of "earlier wins". That becomes an interval scheduling problem: with limit=1, pick a maximum-weight set of non-overlapping pieces.
  4. Load-weighted limit. Shards have a request rate, and the constraint becomes the summed load per key instead of the count. The sweep still works: events carry weights, and "saturated" means load + this shard's load > cap.
  5. Minimum replication. Require every key to have depth at least min_r as well as at most limit. Find the keys whose depth is below min_r with the same sweep. Fill them by extending neighbouring shards, or report that it is impossible without creating new shards.

Interviewer expectations

Mid-level (L4)

  • Asks how ties, empty shards and gaps are handled before coding, and states the "earlier wins, move the start right" rule.
  • Gets a correct solution, even a key-by-key one, and then explains why the key span makes it too slow.
  • Parses and formats "id:start:end" cleanly and handles inclusive ends without off-by-one errors.

Senior (L5)

  • Goes straight to sweeping interval endpoints and explains why "rightmost saturated key + 1" is the smallest legal move.
  • Shows that trimming never opens a hole, so gap filling only handles gaps from the input and never breaks the limit.
  • Writes edge-case tests: a shard inside another, equal starts, adjacent shards (end + 1 == start), and an empty input.

Staff (L6+)

  • Brings up the segment-tree O(n log n) version and the online add/remove design without being asked, including how one early insert spreads changes to all later shards.
  • Treats the output as a migration plan: keys moved, ordering, and keeping both sides readable during a move.
  • Questions the policy itself. "Earlier wins" is simple but can move a lot of data. Discusses other goals (fewest moved keys, weighted load) and their cost.

Sign in to comment, vote or remix.