Rectangle board with remove-at-cell, batch remove and undo

Question: Rectangle board with remove-at-cell, batch remove and undo

Companies: OpenAI

Difficulty: Medium

Problem statement

Build a RectangleBoard that holds axis-aligned rectangles on a grid of width columns and height rows. Cells use zero-indexed (x, y), where x is the column and y is the row. A rectangle is given by its top-left corner (x1, y1) and bottom-right corner (x2, y2). Both corners are inclusive, so (2, 2, 2, 2) covers one cell.

Rectangles may overlap. Where several cover a cell, the one added most recently is on top. Each add gets the next ID, starting at 1. IDs are never reused, even after a rectangle is removed or an add is undone.

The interview comes in parts. Finish each part before you start the next.

Part 1. Implement:

• RectangleBoard(width, height)
• add(x1, y1, x2, y2) -> int: place a rectangle and return its ID. Raise an error if a corner is off the board or the corners are reversed.
• remove(id) -> bool: remove that rectangle. Return False if it is not on the board.
• top_at(x, y) -> int | None: the ID of the top rectangle at the cell, or None.

Part 2. remove_at(x, y) -> int | None: remove the top rectangle at the cell and return its ID, or None if the cell is empty.

Part 3. remove_many(ids) -> int: remove every listed rectangle that is on the board and return how many were removed. Ignore unknown or duplicate IDs.

Part 4. undo() -> bool: reverse the most recent operation that changed the board. Return False if there is nothing to undo. A restored rectangle comes back with its original ID, coordinates and stacking position. A batch from Part 3 is undone in one step. Operations that changed nothing are not undone.

Example. On a 6×4 board:

add(0, 0, 3, 1)   -> 1
add(2, 1, 5, 3)   -> 2
top_at(2, 1)      -> 2
remove_at(2, 1)   -> 2
top_at(2, 1)      -> 1
undo()            -> True    # rectangle 2 is back
top_at(2, 1)      -> 2
add(4, 0, 4, 0)   -> 3
remove_many([1, 3, 7]) -> 2
top_at(0, 0)      -> None
undo()            -> True    # 1 and 3 both return
top_at(4, 0)      -> 3

Constraints: the board is at most 1000×1000, and there are at most 10^4 operations. Say which operation you are optimizing for.

Approach

The question builds up in parts: add and remove by ID, then "remove whatever
is on top at this cell", then batch removal, then undo. None of it needs a
clever algorithm. It tests whether your state stays consistent when
operations are reversed.

Stacking order is the ID order. IDs grow by one with each add and are
never reused. So "added most recently" is the same as "largest ID". The board
does not need a separate z-order list. Restoring a removed rectangle with its
original ID puts it back at its original depth with no extra work.

State.

• rects: dict[int, Rect]: the active rectangles by ID.
• next_id: the next ID to give out. It only goes up. Undoing an add does
  not lower it, because IDs are never reused.
• history: a stack of inverse records, one per operation that changed the
  board. An add pushes ("add", id). Any removal (by ID, at a cell, or a
  batch) pushes ("remove", [(id, rect), ...]). Undo pops one record and
  reverses it. So a batch removal is one record, and one undo brings back the
  whole batch.

Top rectangle at a cell. The brute force scans every active rectangle and
keeps the largest ID that covers (x, y). That is O(R), and it is the right
first answer. It costs nothing on add or remove, and it is easy to show
correct.

The other option is a per-cell index: each cell keeps a max-heap of the IDs
that cover it, and stale IDs are dropped lazily. That makes top_at about
O(log R), but add and remove then cost O(area). With large rectangles on a
small board that trade is bad. State the trade-off and build the scan unless
the interviewer says lookups are much more frequent than writes.

Invariants.

1. Every key in rects is less than next_id, and no ID is in rects
   twice.
2. For each record on history, undoing it gives back exactly the state
   before the operation: the same IDs, coordinates and depth.
3. Operations that change nothing (an unknown ID, an empty cell, a batch
   with no active IDs) push no record. So undo never "spends" itself on a
   no-op.

Code

from dataclasses import dataclass
from typing import Iterable, Optional


@dataclass(frozen=True)
class Rect:
    x1: int
    y1: int
    x2: int
    y2: int

    def covers(self, x: int, y: int) -> bool:
        return self.x1 <= x <= self.x2 and self.y1 <= y <= self.y2


class RectangleBoard:
    def __init__(self, width: int, height: int) -> None:
        if width <= 0 or height <= 0:
            raise ValueError("board must have positive size")
        self.width = width
        self.height = height
        self.rects: dict[int, Rect] = {}
        self.next_id = 1
        # ("add", id) or ("remove", [(id, rect), ...])
        self.history: list[tuple[str, object]] = []

    def _in_bounds(self, x: int, y: int) -> bool:
        return 0 <= x < self.width and 0 <= y < self.height

    # Part 1: add and remove by id
    def add(self, x1: int, y1: int, x2: int, y2: int) -> int:
        if not (self._in_bounds(x1, y1) and self._in_bounds(x2, y2)):
            raise ValueError("rectangle outside the board")
        if x1 > x2 or y1 > y2:
            raise ValueError("corners must be top-left then bottom-right")
        rid = self.next_id
        self.next_id += 1
        self.rects[rid] = Rect(x1, y1, x2, y2)
        self.history.append(("add", rid))
        return rid

    def remove(self, rid: int) -> bool:
        return self.remove_many([rid]) == 1

    def top_at(self, x: int, y: int) -> Optional[int]:
        if not self._in_bounds(x, y):
            return None
        best: Optional[int] = None
        for rid, r in self.rects.items():
            if r.covers(x, y) and (best is None or rid > best):
                best = rid
        return best

    # Part 2: remove the top rectangle at a cell
    def remove_at(self, x: int, y: int) -> Optional[int]:
        rid = self.top_at(x, y)
        if rid is not None:
            self.remove_many([rid])
        return rid

    # Part 3: batch remove, undone as one step
    def remove_many(self, ids: Iterable[int]) -> int:
        removed: list[tuple[int, Rect]] = []
        for rid in dict.fromkeys(ids):  # drop duplicates, keep order
            rect = self.rects.pop(rid, None)
            if rect is not None:
                removed.append((rid, rect))
        if removed:
            self.history.append(("remove", removed))
        return len(removed)

    # Part 4: undo the most recent change
    def undo(self) -> bool:
        if not self.history:
            return False
        kind, data = self.history.pop()
        if kind == "add":
            del self.rects[data]  # next_id stays: ids are never reused
        else:
            for rid, rect in data:
                self.rects[rid] = rect
        return True

    def render(self) -> list[str]:
        """Debug view: top id per cell, '.' when empty."""
        return [
            " ".join(str(self.top_at(x, y) or ".") for x in range(self.width))
            for y in range(self.height)
        ]

Complexity

R is the number of active rectangles, and k is the number of IDs passed to a
batch.

Total space is O(R + size of history). History grows with every change, so
cap it or compact it if the board runs for a long time (see Follow-ups).

Test cases

• 5×4 board. add(0,0,2,2) returns 1. add(1,1,3,3) returns 2.
  top_at(1,1) is 2, top_at(0,0) is 1, top_at(4,3) is None.
• remove_at(1,1) returns 2. Now top_at(1,1) is 1 and top_at(3,3) is
  None.
• undo() after that brings back rectangle 2 with the same ID. top_at(1,1)
  is 2 again: depth is restored, not just presence.
• IDs are never reused: add → 1, undo(), then add → 2.
• remove(99) returns False and pushes no history. A following undo()
  reverses the last real change, not the no-op.
• remove_many([1, 3, 3, 99]) with 1, 2 and 3 active removes 2 rectangles
  and returns 2. One undo() restores both 1 and 3.
• remove_at on an empty cell returns None and pushes no history.
• A one-cell rectangle add(2,2,2,2) covers only (2,2).
• Out of bounds: add(0,0,5,1) on a 5-wide board raises ValueError, and
  top_at(-1, 0) returns None.
• Reversed corners: add(3,3,1,1) raises ValueError.
• undo() on a new board returns False.
• A full undo chain: add 1, add 2, remove_many([1, 2]), then three undos
  leaves an empty board, and next_id is still 3.

Follow-ups

1. Redo. Keep a second stack. undo moves its record onto the redo
   stack. redo re-applies it: for an add, put the same ID and rectangle
   back (store the rectangle in the record), and for a removal, pop the IDs
   again. Any new change clears the redo stack.
2. Move or resize a rectangle. Decide whether a move keeps the
   rectangle's depth or brings it to the front. Keeping the depth is a
   ("move", id, old_rect) record. Bringing it to the front changes its
   identity under "ID equals depth", so either give it a new ID or add an
   explicit z counter separate from the ID. The z counter is cleaner and
   also supports "bring to front" and "send to back".
3. Fast top_at with many rectangles. For a small board with many
   lookups, keep a max-heap of IDs per cell and pop dead IDs lazily on read.
   For a large or unbounded board, bucket rectangles into a coarse grid of
   tiles, and scan only the tile that holds the cell, in descending ID order.
   For something more exact, use an interval tree or segment tree on x with
   y filtering. Name the write-cost trade-off.
4. Count or list what is visible. "Which rectangles show at least one
   cell?" or "what area does each one show?" Paint from the bottom to the top
   into a grid, O(sum of areas). Or paint from the top down and skip cells
   that are already painted. With coordinate compression this works when the
   board is huge but the rectangles are few.
5. Unbounded history. Cap the undo depth with a deque, or take
   periodic snapshots and keep only the records after them. If the board is
   shared by several users, each user's undo should reverse only that user's
   own operations. That is a step toward operational transforms, so say how
   you would scope it.

Interviewer expectations

Mid-level (L4)

• Gets add, remove and top_at working, with IDs that start at 1 and are
  never reused, and validates corners and bounds.
• Implements undo for a single add or remove, and handles undo on an empty
  history.
• Tests overlap and the no-op cases by hand.

Senior (L5)

• Sees that "most recent" equals "largest ID", so a restored rectangle keeps
  its depth without a separate z-order.
• Uses one inverse-record stack for every change, so a batch removal undoes
  as one step, and no-ops push no records.
• Explains the trade-off between scanning and a per-cell index, and picks
  one on purpose.

Staff (L6+)

• Builds the parts in increments without rewrites, and the history design
  carries over to redo and move.
• Separates depth from ID once "bring to front" or move appears, and
  explains why.
• Talks about scale and lifecycle: indexes for large boards, bounding
  history, and how undo would change with several users sharing one board.