Implement-Then-Extend Problems

The xAI coding round has a shape, and the shape is the test. You are handed one component to implement and get it running, and then the interviewer changes the requirements while you watch. A feature is bolted on, a constraint is tightened, or a case you did not model becomes the case. Treat the opening prompt as the base you will extend. A candidate who reached a passing first version through small methods and cleanly separated state absorbs the twist in a few lines, while one who hard-coded an assumption starts over under a shrinking clock. This lesson works three of the recurring components in full, in the order the round unfolds: the simple core first, then the live extension it is built to take.

The xAI angle

Choose the base structure for the follow-up you cannot see yet, because at xAI it is almost always coming. Each problem below has a naive first draft that passes the opening prompt and then collapses on the extension. A single "in a transaction" boolean dies on nesting, an iterator that only exposes next and has_next cannot be checkpointed, and a rate limiter with one counter cannot enforce a second rule. In every case the fix lies in the structure you pick: a stack of overlays, an explicit position, or a pair of sliding windows makes the extension one more push. So build the ordinary version out of parts that come apart, and say out loud that you expect the requirement to move.

The implement-then-extend loop

Every problem in this round yields to the same discipline. First, scope the base with a clarifying question or two, then implement the smallest correct version out of small, independent operations, testing each against a case you name before the interviewer does. Second, when the extension lands, resist rewriting: ask what the new rule changes about the state you already hold, and reach for the smallest structural change that carries it. The usual failure is racing to a tangled first version that happens to pass, which has nowhere to grow when "now make it nestable" or "now add a second limit" arrives. Clean and extensible beats fast and fragile here, and the interviewer is grading the extensible part.

Key-value store: growing nested transactions

The most-cited xAI problem is an in-memory key-value store, and it is the cleanest illustration of the whole loop. The base is trivial: set a key, get a key, delete a key, backed by one dictionary. Then the interviewer adds a transaction layer with begin, commit, and rollback, so a block of writes can be discarded as a group. Then comes the constraint that separates a passing answer from a strong one. Nested transactions mean a begin inside an open transaction has to commit into its parent or roll back on its own, independently of anything above it. The structure you pick for the transaction layer decides whether nesting costs you three lines or a teardown.

The draft that dies here keeps a single boolean, "am I in a transaction," and one side buffer of pending writes, which handles exactly one level. The structure that grows is a stack of overlay dictionaries: committed data sits in a base dictionary, and each open transaction pushes its own overlay on top. A read walks the stack from the newest overlay down to the base and returns the first value it finds. A write lands in the topmost overlay, or in the base when no transaction is open, and commit folds the top overlay into the one beneath it while rollback discards it. Nesting is free, because a begin inside a transaction is just one more push. The one subtlety is delete inside a transaction: since a rollback must bring the key back, it writes a tombstone marker into that layer, which shadows the value below until it resolves.

1

In-Memory Key-Value Store With Nested Transactions

MediumFrequently asked

Problem

Implement a KVStore supporting reads, writes, deletes, and nested transactions:

set(key, value) stores a value. get(key) returns the current value for a key, or null if the key is not set. delete(key) removes a key.

begin() opens a new transaction. commit() closes the innermost open transaction, folding its changes into the enclosing transaction if one is open, or into the committed store otherwise. rollback() closes the innermost open transaction and discards all of its changes.

Transactions nest: a begin inside an open transaction creates a child whose commit or rollback resolves independently of its parent. A delete performed inside a transaction must be undoable by a rollback.

Input: set("a","1"); begin(); set("a","2"); get("a"); rollback(); get("a")

Output: "2", "1"

Inside the transaction the value is "2". Rolling back discards the transaction, so the committed "1" is restored.

Input: set("k","base"); begin(); set("k","outer"); begin(); set("k","inner"); commit(); get("k"); rollback(); get("k")

Output: "inner", "base"

The inner commit folds "inner" into the outer transaction, so get sees "inner". Rolling back the outer transaction discards everything, restoring "base".

How this uses the pattern

Keep committed data in a base dictionary and one overlay dictionary per open transaction on a stack. Read newest-overlay-first down to the base. Write to the top overlay. Commit folds the top overlay down one level; rollback pops it. Model an in-transaction delete as a tombstone in the overlay so rollback can restore the shadowed value.

In-Memory Key-Value Store With Nested Transactions

Python

Focus Mode

Loading...

OutputSubmissions

Run

Submit

Click "Run" to execute your code

Test Cases

Input: [["KVStore",[]],["set",["a","1"]],["set",["b","2"]],["get",["a"]],["delete",["a"]],["get",["a"]],["get",["b"]]]Expected: [null,null,null,"1",null,null,"2"]

Input: [["KVStore",[]],["set",["a","1"]],["get",["a"]],["begin",[]],["set",["a","2"]],["get",["a"]],["begin",[]],["set",["a","3"]],["delete",["a"]],["get",["a"]],["rollback",[]],["get",["a"]],["commit",[]],["get",["a"]],["get",["b"]]]Expected: [null,null,"1",null,null,"2",null,null,null,null,null,"2",null,"2",null]

Input: [["KVStore",[]],["set",["k","base"]],["begin",[]],["set",["k","outer"]],["begin",[]],["set",["k","inner"]],["commit",[]],["get",["k"]],["rollback",[]],["get",["k"]]]Expected: [null,null,null,null,null,null,null,"inner",null,"base"]

Input: [["KVStore",[]],["set",["x","10"]],["begin",[]],["delete",["x"]],["get",["x"]],["commit",[]],["get",["x"]]]Expected: [null,null,null,null,null,null,null]

Need a hint? (3 available)

The xAI angle

The key-value store is the signature xAI problem because of the nesting extension. A first draft with a single boolean transaction flag passes the begin-commit-rollback prompt. It has to be rewritten the moment a begin lands inside another begin. A stack of overlays passes both without a rewrite, because nesting is one more push and rollback is one pop. Interviewers watch which structure you reach for before the nesting is even announced. Choosing the stack up front shows you expect the extension.

Take it further

Add a time-to-live to each key. Suppose set takes an expiry timestamp and get takes the current time, so a key read after its expiry returns null as if unset. Store the expiry alongside the value in each overlay entry. Have the newest-overlay-first read skip any entry whose expiry has passed, treating it like the tombstone case. The transaction machinery does not change. Expiry rides along inside the same overlay values that commit and rollback already move, so the feature is a check at read time with no new layer of state.

Show solution & walkthrough

Stateful iterator: checkpoint and resume

The second recurring shape is an iterator, and checkpointing is the extension that makes it an xAI problem. The base is an ordinary iterator over run-length encoded pairs, where ["a", 2] means the value "a" repeated twice, exposed through next and has_next. The interviewer then adds get_state and set_state, so a caller can snapshot the iterator mid-stream, keep advancing, and later restore and resume from the exact snapshot point. The base draft that stalls hides its progress inside a Python generator or a decompressed list it cannot address. The version that extends holds its read position as plain fields, so a checkpoint is just copying those fields out and a restore is copying them back.

2

Stateful Iterator With Checkpoint and Resume

MediumCommonly asked

Problem

Implement an RLEIterator over a run-length encoded sequence. The constructor takes a list of [value, count] runs, so [["a",2],["b",1]] represents the stream a, a, b.

next() returns the next value in the stream, or null when the stream is exhausted. has_next() returns whether any values remain.

Extend it to be checkpointable: get_state() returns a snapshot of the current position, and set_state(state) restores the iterator to a previously returned snapshot, so iteration resumes from exactly that point.

Input: RLEIterator([["a",2],["b",1],["c",3]]); next(); next(); s = get_state(); next(); next(); set_state(s); next()

Output: "a", "a", "b", "c", "b"

After two calls the snapshot sits at the start of the "b" run. Two more calls advance to "b" then "c". Restoring the snapshot rewinds to "b", so the next call returns "b" again.

Input: RLEIterator([["x",1]]); has_next(); next(); has_next(); next()

Output: true, "x", false, null

One value in the stream. After it is consumed, has_next is false and a further next returns null.

How this uses the pattern

Hold the read position as two plain fields: the index of the current run and how many of it you have consumed. Each next call reads the current run and advances those fields, rolling to the next run when the current one is used up. The get_state call returns a copy of the two fields, and set_state writes them back.

Stateful Iterator With Checkpoint and Resume

Python

Focus Mode

Loading...

OutputSubmissions

Run

Submit

Click "Run" to execute your code

Test Cases

Input: [["RLEIterator",[[["a",2],["b",1],["c",3]]]],["next",[]],["next",[]],["get_state",[]],["next",[]],["next",[]],["set_state",[[1,0]]],["next",[]],["next",[]],["has_next",[]]]Expected: [null,"a","a",[1,0],"b","c",null,"b","c",true]

Input: [["RLEIterator",[[["x",1]]]],["has_next",[]],["next",[]],["has_next",[]],["next",[]]]Expected: [null,true,"x",false,null]

Input: [["RLEIterator",[[["p",3]]]],["next",[]],["get_state",[]],["next",[]],["next",[]],["has_next",[]],["set_state",[[0,1]]],["next",[]],["has_next",[]]]Expected: [null,"p",[0,1],"p","p",false,null,"p",true]

Need a hint? (3 available)

The xAI angle

The iterator is an xAI favorite because get_state and set_state expose whether your progress lives somewhere you can address. A generator-based or pre-decompressed base looks clean, then cannot implement checkpointing at all, which forces a rewrite. Holding the position as explicit fields makes the extension trivial. A snapshot is those fields copied out, and a restore is them copied in. The interviewer is checking that your base kept its state in the open, outside the language machinery.

Take it further

Support set_state to any earlier snapshot, including one taken across a reset. The state is just the run index and the offset into that run. A snapshot stays valid no matter how far you have advanced past it, so you can hold several snapshots and jump between them freely. Guard set_state by clamping a restored index that points past the final run to the exhausted position. A stale snapshot then degrades to a clean has_next of false instead of an out-of-bounds read. The stored state never has to grow beyond those two integers.

Show solution & walkthrough

Multi-level rate limiter: one key, then compound rules

The third shape is a rate limiter, and its extension is a second rule. The base admits or denies a request against one limit, at most some number of requests per key within a sliding time window, decided from a stream of timestamps in nondecreasing order. The interviewer then adds a second, wider limit that applies at the same time, a global cap across every key, so a request is admitted only when it clears both budgets. The draft that stalls keeps one counter and has no room for a second rule. The draft that extends keeps one sliding window per rule and admits only when every window has room, so adding a level means adding a window and one more clause to the and.

3

Multi-Level Rate Limiter Over a Timestamp Stream

MediumCommonly asked

Problem

Implement a RateLimiter that decides, for a stream of requests arriving in nondecreasing timestamp order, whether each request is admitted or denied.

The constructor takes per_key_limit, global_limit, and window. allow(key, timestamp) returns true if the request is admitted and false if it is denied.

A request is admitted only when both limits have room. That means at most per_key_limit admitted requests for that key within the last window units of time, and at most global_limit admitted requests across all keys within the same window. A window is the half-open span (timestamp - window, timestamp], so a request made exactly window units after an earlier one no longer counts that earlier one. Only admitted requests consume budget; a denied request records nothing.

Input: RateLimiter(2,3,10); allow("a",1); allow("a",2); allow("a",3); allow("b",4); allow("b",5); allow("a",11)

Output: true, true, false, true, false, true

Key "a" fills its two slots at t=1 and t=2, so t=3 is denied. "b" at t=4 fills the third global slot. "b" at t=5 is denied because the global limit of 3 is full. By t=11 the t=1 event has left the window, freeing both an "a" slot and a global slot, so "a" at t=11 is admitted.

Input: RateLimiter(5,2,100); allow("a",1); allow("b",1); allow("c",1); allow("a",2)

Output: true, true, false, false

The per-key limit is loose, but the global limit of 2 fills at t=1 across keys "a" and "b". The third and fourth requests are denied even though each key is under its own budget.

How this uses the pattern

Keep one sliding window of admitted timestamps per key and one shared window for the global count. On each request, evict timestamps at or before timestamp minus window from both. Admit only when both windows are below their limits, and on admit append the timestamp to both.

Multi-Level Rate Limiter Over a Timestamp Stream

Python

Focus Mode

Loading...

OutputSubmissions

Run

Submit

Click "Run" to execute your code

Test Cases

Input: [["RateLimiter",[2,3,10]],["allow",["a",1]],["allow",["a",2]],["allow",["a",3]],["allow",["b",4]],["allow",["b",5]],["allow",["a",11]],["allow",["b",12]]]Expected: [null,true,true,false,true,false,true,true]

Input: [["RateLimiter",[1,5,1]],["allow",["x",1]],["allow",["x",1]],["allow",["x",2]],["allow",["y",2]]]Expected: [null,true,false,true,true]

Input: [["RateLimiter",[5,2,100]],["allow",["a",1]],["allow",["b",1]],["allow",["c",1]],["allow",["a",2]]]Expected: [null,true,true,false,false]

Need a hint? (3 available)

The xAI angle

A rate limiter is an xAI staple because the extension is a clean test of separated state. A single counter passes the one-limit prompt but cannot express a second rule, so it forces a rewrite the moment a global cap is added. One sliding window per rule makes the compound case an and across windows. A third level later is a third window, with no restructuring. The interviewer is watching whether your base kept each rule's history in its own structure. Merging them into one number you cannot pull apart is the failure.

Take it further

Add a per-endpoint limit as a third level. Suppose each request also names a URL, with a cap per URL within the window on top of the per-key and global caps. Each rule already owns its own sliding window, so the change is a third dictionary of windows keyed by URL and a third clause in the admit check. On admit, the timestamp goes into all three windows. No existing rule's bookkeeping changes, which is the payoff of keeping the windows separate from the first line.

Show solution & walkthrough

What these three share

Look across the three and the same choice wins every time. The store, the iterator, and the limiter each have a naive base that passes the opening prompt and cannot take the extension. Each also has a base built from separated state that takes the extension in a few lines. That means a stack of overlays over a boolean, explicit position fields over a generator, and one window per rule over one counter. Rather than memorizing these three solutions, train the reflex to ask, before you write the base, where the requirement is most likely to move and which structure keeps that move cheap.

Common questions