Indexed sequence with checkpoints

Question: Indexed sequence with checkpoints

Companies: OpenAI

Difficulty: Medium

Problem statement

Build a growable array that can take cheap snapshots. The spec arrives in three parts; finish and test each part before the next.

Part 1. append(value) -> index, pop() -> value, get(index), set(index, value) and len(). Indexes run from 0 to len - 1; anything else raises IndexError (no negative indexing).

Part 2. checkpoint() -> id returns ids 0, 1, 2, …. get_at(id, index) returns the value index had when that checkpoint was taken, and len_at(id) the length then. Writes made after checkpoint() returns are not part of it. An index that did not exist at the checkpoint raises IndexError; an unknown checkpoint raises KeyError. A checkpoint must not copy the array: assume 10^5 items and 10^4 checkpoints.

Part 3. rollback(id) makes the live array equal checkpoint id again. Checkpoint id stays valid; every later checkpoint is forgotten, and the next checkpoint() returns id + 1. Rollback must not scan the whole array.

class IndexedSequence:
    def __len__(self) -> int: ...
    def append(self, value) -> int: ...
    def pop(self): ...
    def get(self, index: int): ...
    def set(self, index: int, value) -> None: ...
    def checkpoint(self) -> int: ...
    def get_at(self, checkpoint: int, index: int): ...
    def len_at(self, checkpoint: int) -> int: ...
    def rollback(self, checkpoint: int) -> None: ...

Example

s.append(10); s.append(20)
c0 = s.checkpoint()        # 0
s.set(0, 11); s.append(30)
s.get_at(c0, 0)  -> 10
s.get_at(c0, 2)  -> IndexError   (index 2 did not exist at c0)
s.rollback(c0)
len(s) -> 2, s.get(0) -> 10
s.checkpoint()   -> 1

Approach

Three levels, each one changing what "fast" has to mean:

1. A plain list does append, pop, get, set in O(1).
2. checkpoint() must be cheap even when the list holds a million items, and get_at(c, i) must read the value index i had at checkpoint c.
3. rollback(c) makes the live list equal checkpoint c again and forgets every later checkpoint.

Brute force. checkpoint() copies the whole list. Reads at a checkpoint are O(1), but a checkpoint is O(n) time and memory. With n = 10^5 and 10^5 checkpoints that is 10^10 element copies: say this, then move on.

Version-tagged history per index. Keep a version counter that equals the number of checkpoints taken so far. Every index holds its own history: a list of (version, value) pairs in increasing version order. A write at the current version either appends a new pair, or overwrites the last pair if it already has the current version. Then:

• checkpoint() returns the current version and increments it: O(1), no copying.
• Checkpoint c sees, for each index, the newest pair with version <= c. The versions in one index's history are sorted, so that is one bisect: O(log h), where h is that index's history length.
• The length changes over time too (append, pop), so it gets the same treatment: one history of (version, length) pairs. get_at(c, i) first checks i < len_at(c).

Overwriting within a version is what keeps memory at O(n + number of (index, checkpoint) pairs that actually changed), not O(writes). A million set(0, ...) calls between two checkpoints leave one pair.

Rollback with an undo log. Rolling back to c means dropping every pair with version greater than c, in every index history and in the length history. Scanning all n indexes is O(n). Instead, append a log entry every time a new pair is created (not on an overwrite): ("set", i) or ("len",). Versions only grow while you write, so the log is sorted by the version of the pair each entry created. Rollback pops log entries from the end while their version is greater than c, and pops the matching pair. Then set the version counter to c + 1.

Each pair is created once and popped at most once, so rollback is O(1) amortized per pair it removes. After the rollback the log is still sorted (everything newer than c is gone), so later writes and rollbacks keep working.

Definitions to say out loud before coding. These are the questions the interviewer is waiting for you to ask:

• Is the state at checkpoint c taken before or after writes made right after checkpoint() returned? After the call, writes belong to the next version, so c sees only what happened before the call.
• After rollback(c), is c still valid? Yes. Are later ids reused? Yes, the next checkpoint() returns c + 1.
• Can get_at read an index that did not exist at that checkpoint? No: IndexError, even if the index exists now.
• Unknown checkpoint id: KeyError.

Code

from __future__ import annotations

from bisect import bisect_right


class IndexedSequence:
    """A growable array with cheap checkpoints.

    Every write is tagged with the current version, the number of checkpoints
    taken so far. Checkpoint c sees every write tagged <= c.
    """

    def __init__(self) -> None:
        self._version = 0
        self._vals: list = []      # per index: values, oldest first
        self._vers: list = []      # per index: version of each value
        self._len_vers: list = [0]  # versions at which the length changed
        self._lens: list = [0]      # the length from that version on
        # Undo log of every pair created, in order: ("set", index) or
        # ("len",). Lets rollback drop newer pairs without scanning.
        self._log: list = []

    # -- level 1 -----------------------------------------------------------
    def __len__(self) -> int:
        return self._lens[-1]

    def append(self, value) -> int:
        i = len(self)
        if i == len(self._vals):
            self._vals.append([])
            self._vers.append([])
        self._set_len(i + 1)
        self._write(i, value)
        return i

    def pop(self):
        n = len(self)
        if n == 0:
            raise IndexError("pop from empty sequence")
        value = self._vals[n - 1][-1]
        self._set_len(n - 1)
        return value

    def set(self, index: int, value) -> None:
        self._check(index, len(self))
        self._write(index, value)

    def get(self, index: int):
        self._check(index, len(self))
        return self._vals[index][-1]

    # -- level 2 -----------------------------------------------------------
    def checkpoint(self) -> int:
        c = self._version
        self._version += 1
        return c

    def get_at(self, checkpoint: int, index: int):
        self._check_checkpoint(checkpoint)
        k = bisect_right(self._len_vers, checkpoint) - 1
        self._check(index, self._lens[k])
        j = bisect_right(self._vers[index], checkpoint) - 1
        return self._vals[index][j]

    def len_at(self, checkpoint: int) -> int:
        self._check_checkpoint(checkpoint)
        return self._lens[bisect_right(self._len_vers, checkpoint) - 1]

    # -- level 3 -----------------------------------------------------------
    def rollback(self, checkpoint: int) -> None:
        """Make the live state equal checkpoint `checkpoint`. Later checkpoints
        are forgotten; `checkpoint` itself stays valid."""
        self._check_checkpoint(checkpoint)
        while self._log and self._log_version(self._log[-1]) > checkpoint:
            entry = self._log.pop()
            if entry[0] == "set":
                self._vals[entry[1]].pop()
                self._vers[entry[1]].pop()
            else:
                self._lens.pop()
                self._len_vers.pop()
        self._version = checkpoint + 1

    # -- helpers -----------------------------------------------------------
    def _write(self, index: int, value) -> None:
        vers, vals = self._vers[index], self._vals[index]
        if vers and vers[-1] == self._version:
            vals[-1] = value  # same version: overwrite, history unchanged
            return
        vers.append(self._version)
        vals.append(value)
        self._log.append(("set", index))

    def _set_len(self, n: int) -> None:
        if self._len_vers[-1] == self._version:
            self._lens[-1] = n
            return
        self._len_vers.append(self._version)
        self._lens.append(n)
        self._log.append(("len",))

    def _log_version(self, entry) -> int:
        if entry[0] == "set":
            return self._vers[entry[1]][-1]
        return self._len_vers[-1]

    def _check(self, index: int, length: int) -> None:
        if not 0 <= index < length:
            raise IndexError(index)

    def _check_checkpoint(self, checkpoint: int) -> None:
        if not 0 <= checkpoint < self._version:
            raise KeyError(checkpoint)

Two details that are easy to get wrong:

• pop keeps the popped value in the history. An older checkpoint can still read it. Only the length history changes. If the same slot is appended again in a later version, it gets a new pair; in the same version, the pair is overwritten.
• The log entry finds its own pair. A ("set", i) entry is always the newest pair of index i at the moment it is popped, because pairs for i were created in log order. So the entry does not need to store the version.

Complexity

Let n be the current length, w the number of writes, and p the number of (index, version) pairs that exist (p ≤ w, and usually far smaller).

Worked numbers: 10^5 items, 10^4 checkpoints, 100 sets between checkpoints. Brute force copies 10^5 × 10^4 = 10^9 slots (about 8 GB of pointers). Version tagging stores at most 100 × 10^4 = 10^6 pairs plus the 10^5 originals, about 1.1 × 10^6 pairs, and get_at bisects a history that averages ten pairs per touched index.

Tests

from sequence import IndexedSequence


def raises(exc, fn, *args):
    try:
        fn(*args)
    except exc:
        return True
    return False


def test_append_get_set():
    s = IndexedSequence()
    assert s.append("a") == 0
    assert s.append("b") == 1
    s.set(0, "z")
    assert [s.get(0), s.get(1)] == ["z", "b"]
    assert len(s) == 2


def test_out_of_range():
    s = IndexedSequence()
    s.append(1)
    assert raises(IndexError, s.get, 1)
    assert raises(IndexError, s.set, -1, 5)
    assert s.pop() == 1
    assert raises(IndexError, s.pop)

from sequence import IndexedSequence


def raises(exc, fn, *args):
    try:
        fn(*args)
    except exc:
        return True
    return False


def test_checkpoint_reads_old_values():
    s = IndexedSequence()
    s.append(10)
    s.append(20)
    c0 = s.checkpoint()
    s.set(0, 11)
    s.append(30)
    c1 = s.checkpoint()
    s.set(0, 12)
    assert (c0, c1) == (0, 1)
    assert s.get_at(c0, 0) == 10
    assert s.get_at(c1, 0) == 11
    assert s.get(0) == 12
    assert s.len_at(c0) == 2 and s.len_at(c1) == 3
    assert raises(IndexError, s.get_at, c0, 2)  # index 2 did not exist yet
    assert raises(KeyError, s.get_at, 5, 0)     # no such checkpoint


def test_many_writes_between_checkpoints():
    s = IndexedSequence()
    s.append(0)
    for v in range(1000):
        s.set(0, v)
    c = s.checkpoint()
    s.set(0, -1)
    assert s.get_at(c, 0) == 999 and s.get(0) == -1


def test_pop_then_append_keeps_old_checkpoint():
    s = IndexedSequence()
    s.append("x")
    c = s.checkpoint()
    s.pop()
    s.append("y")
    assert s.get_at(c, 0) == "x" and s.get(0) == "y"

from sequence import IndexedSequence


def raises(exc, fn, *args):
    try:
        fn(*args)
    except exc:
        return True
    return False


def test_rollback_restores_values_and_length():
    s = IndexedSequence()
    s.append(1)
    s.append(2)
    c0 = s.checkpoint()
    s.set(1, 20)
    s.append(3)
    c1 = s.checkpoint()
    s.append(4)
    s.rollback(c0)
    assert len(s) == 2 and s.get(1) == 2
    assert raises(KeyError, s.get_at, c1, 0)  # later checkpoints are gone
    assert s.get_at(c0, 1) == 2               # c0 itself still works
    c = s.checkpoint()
    assert c == c0 + 1                        # ids are reused after rollback
    s.append(99)
    assert s.get(2) == 99 and s.get_at(c, 1) == 2


def test_rollback_twice_and_write_in_between():
    s = IndexedSequence()
    s.append("a")
    c0 = s.checkpoint()
    s.set(0, "b")
    s.rollback(c0)
    s.set(0, "c")
    s.rollback(c0)
    assert s.get(0) == "a"


def test_against_a_brute_force_model():
    import random

    rng = random.Random(7)
    s = IndexedSequence()
    live = []
    snaps = []  # full copies, list index = checkpoint id
    for _ in range(3000):
        op = rng.random()
        if op < 0.3 or not live:
            v = rng.randint(0, 99)
            s.append(v)
            live.append(v)
        elif op < 0.55:
            i = rng.randrange(len(live))
            v = rng.randint(0, 99)
            s.set(i, v)
            live[i] = v
        elif op < 0.65:
            assert s.pop() == live.pop()
        elif op < 0.8:
            assert s.checkpoint() == len(snaps)
            snaps.append(list(live))
        elif op < 0.85 and snaps:
            c = rng.randrange(len(snaps))
            s.rollback(c)
            del snaps[c + 1:]
            live = list(snaps[c])
        elif snaps:
            c = rng.randrange(len(snaps))
            assert s.len_at(c) == len(snaps[c])
            if snaps[c]:
                i = rng.randrange(len(snaps[c]))
                assert s.get_at(c, i) == snaps[c][i]
        assert len(s) == len(live)
        assert all(s.get(i) == live[i] for i in range(len(live)))

The last test is the one to write in the interview if you have five spare minutes: the brute force (full copies) is the oracle, and 3,000 random operations check every level at once.

Follow-ups

• "Release checkpoints nobody needs." release(c) means no one will read c again. Pairs only exist to answer reads at some live checkpoint, so for each index you can merge a pair into the next one when no live checkpoint falls between their versions. Do it lazily: keep the live checkpoints in a sorted set, and compact one index's history when you next write to it. Rollback targets must stay live, so rollback to a released checkpoint raises.
• "Rollback, then branch." Today a rollback destroys the newer checkpoints. To keep them (like Git branches), stop truncating and make the version graph a tree: each checkpoint has a parent, and a read at c walks the path from c to the root. That is a persistent data structure; a persistent segment tree or a fat-node array with path copying gives O(log n) reads and writes per version.
• "Insert and delete in the middle." Index-tagged histories break, because an insert shifts every later index. Use a persistent balanced tree (a rope or an implicit treap with path copying). Each write copies O(log n) nodes, a checkpoint is a pointer to the root, and rollback is "make that root current".
• "Make it durable." The undo log is almost a write-ahead log already. Write each pair to disk as (version, index, value) before applying it, and write checkpoint and rollback records too. Recovery replays the log; a periodic snapshot of the live list (plus the length history) bounds replay time.
• "Concurrent readers." Readers at an old checkpoint never conflict with writers at the current version, which is the point of MVCC. With one writer and a lock-free append-only history per index, readers need no lock as long as list appends are atomic, which holds in CPython but should not be relied on across languages.

Interviewer expectations

Mid-level (L4)

• Level 1 correct with bounds checks; a working level 2, even if checkpoints start as full copies.
• Moves to per-index history with binary search when the interviewer asks about memory, and handles "index did not exist at that checkpoint".
• Tests each level before moving on.

Senior (L5)

• Goes straight to version tagging with overwrite-within-version and explains why memory is O(changed pairs), not O(writes).
• Asks the semantics questions up front (what c sees, id reuse, rollback validity) and encodes the answers in tests.
• Implements rollback without scanning every index (an undo log or equivalent) and states the amortized cost.

Staff (L6+)

• Writes the brute-force model test and uses it to find their own bugs.
• Connects the design to MVCC and snapshot isolation, and discusses releasing checkpoints, branching histories and persistent trees with their trade-offs.
• Explains durability (the undo log as a WAL plus snapshots) and what changes once inserts in the middle are allowed.