Select table cells and set a color

Clarify first. Pin these down before you code:

• What a select* call does. It only remembers the selection. It does not change any color. A new select* call replaces the old selection.
• setColor with nothing selected. Ask. A common answer is "do nothing". Raising an error is also fine; say which one you chose.
• Bad indices. selectRow(r) or selectCell(r, c) outside the table should raise an error and keep the old selection.
• Last write wins. If you color row 2 red and then cell (2, 3) blue, cell (2, 3) is blue and the rest of row 2 stays red. If you then color the whole table green, every cell is green.
• Which operation is hot. Tables can be large. Ask whether setColor on the whole table must be fast. (It must, or the interviewer will push you there.)

Brute force: a full grid. Keep an R x C array of colors. setColor writes every cell in the selection. getColor reads one cell.

• setColor on the table is O(R * C). On a row it is O(C). On a cell it is O(1).
• getColor is O(1).
• Memory is O(R * C) even when nothing is colored.

This is correct and a fine place to start. The interviewer will then ask for a large table (for example a million rows) where setColor on the whole table must not touch every cell.

Optimal: store each paint at its own level. A selection is one of three shapes: the table, one row, or one cell. Store a paint where it happened, not in every cell it covers:

• table_color: the last color set on the whole table.
• row_color[r]: the last color set on row r, if it is newer than the table color.
• cell_color[r][c]: the last color set on cell (r, c), if it is newer than its row color and the table color.

getColor(r, c) checks the most specific level first: the cell, then the row, then the table.

Why the most specific level wins. Keep this invariant:

Every stored entry is newer than every coarser entry that covers it.

setColor keeps it true by deleting the finer entries it overwrites:

• Table paint: set table_color, then clear all row and cell entries. They are all older now.
• Row paint on row r: set row_color[r], then drop cell_color[r]. Those cells are older now.
• Cell paint: set cell_color[r][c]. It is the newest thing that covers that cell.

Because of the invariant, the first entry getColor finds is the newest paint for that cell. So there is no need to store timestamps for these three shapes.

Cost of the clears. Replacing the row and cell maps with new empty maps is O(1) for the caller. Each entry is created once and freed once, so the total cleanup work is never more than the number of setColor calls made so far (amortized O(1)).

Alternative: timestamps. Give every paint an increasing counter and keep (stamp, color) at each level. getColor returns the color with the largest stamp among the cell, its row and the table. This needs no clearing and works for shapes that do not nest, like columns (see Follow-ups). Without clearing, memory grows with the number of paints.

from __future__ import annotations


class Table:
    """A table whose cells can be colored by table, row, or cell selection.

    Invariant: every stored entry is newer than every coarser entry that
    covers it, so getColor can return the most specific entry it finds.
    """

    def __init__(self, num_rows: int, num_columns: int) -> None:
        if num_rows <= 0 or num_columns <= 0:
            raise ValueError("table must have at least one row and one column")
        self.num_rows = num_rows
        self.num_columns = num_columns
        self._table_color = ""
        self._row_color: dict[int, str] = {}
        self._cell_color: dict[int, dict[int, str]] = {}
        # ("table",), ("row", r), ("cell", r, c), or None
        self._selection: tuple[int | str, ...] | None = None

    # --- selection -------------------------------------------------------
    def selectTable(self) -> None:
        self._selection = ("table",)

    def selectRow(self, row: int) -> None:
        self._check_row(row)
        self._selection = ("row", row)

    def selectCell(self, row: int, col: int) -> None:
        self._check_row(row)
        self._check_col(col)
        self._selection = ("cell", row, col)

    # --- coloring --------------------------------------------------------
    def setColor(self, color: str) -> None:
        sel = self._selection
        if sel is None:
            return  # nothing selected: no-op
        kind = sel[0]
        if kind == "table":
            self._table_color = color
            self._row_color = {}  # older than the new table color
            self._cell_color = {}
        elif kind == "row":
            row = sel[1]
            self._row_color[row] = color
            self._cell_color.pop(row, None)  # cells in this row are older
        else:
            _, row, col = sel
            self._cell_color.setdefault(row, {})[col] = color

    def getColor(self, row: int, col: int) -> str:
        self._check_row(row)
        self._check_col(col)
        cells = self._cell_color.get(row)
        if cells is not None and col in cells:
            return cells[col]
        if row in self._row_color:
            return self._row_color[row]
        return self._table_color

    # --- helpers ---------------------------------------------------------
    def _check_row(self, row: int) -> None:
        if not 0 <= row < self.num_rows:
            raise IndexError(f"row {row} out of range")

    def _check_col(self, col: int) -> None:
        if not 0 <= col < self.num_columns:
            raise IndexError(f"column {col} out of range")


if __name__ == "__main__":
    t = Table(3, 4)
    assert t.getColor(0, 0) == ""
    t.setColor("red")  # nothing selected
    assert t.getColor(0, 0) == ""

    t.selectRow(1)
    t.setColor("red")
    t.selectCell(1, 2)
    t.setColor("blue")
    assert t.getColor(1, 2) == "blue"
    assert t.getColor(1, 0) == "red"
    assert t.getColor(0, 0) == ""

    t.selectRow(1)
    t.setColor("yellow")  # overwrites the blue cell
    assert t.getColor(1, 2) == "yellow"

    t.selectTable()
    t.setColor("green")
    assert all(t.getColor(r, c) == "green" for r in range(3) for c in range(4))

    t.selectCell(2, 3)
    t.setColor("pink")
    t.setColor("black")  # selection persists
    assert t.getColor(2, 3) == "black"
    assert t.getColor(2, 2) == "green"
    print("ok")

Let R be rows and C be columns.

Space is O(P), where P is the number of row and cell entries still stored. It is never more than R + R * C, and in practice it is far smaller. The brute-force grid uses O(R * C) space and O(R * C) time for a table paint.

• New Table(3, 4): getColor(0, 0) returns "".
• setColor("red") before any select* call: no cell changes.
• selectRow(1), setColor("red"): getColor(1, 0) is "red", getColor(0, 0) is "".
• Then selectCell(1, 2), setColor("blue"): getColor(1, 2) is "blue", getColor(1, 3) is still "red".
• Then selectRow(1), setColor("yellow"): getColor(1, 2) is "yellow" (the row paint is newer than the cell paint).
• Then selectTable(), setColor("green"): every cell is "green".
• selectCell(2, 3), then setColor("pink"), then setColor("black"): the selection stays, so getColor(2, 3) is "black".
• Table(1, 1): selectRow(0) and selectCell(0, 0) both cover the single cell.
• selectRow(3) on a 3-row table raises IndexError, and the old selection is kept.
• getColor(0, 4) on a 4-column table raises IndexError.
• Large table, Table(10**6, 10**6): selectTable(), setColor("x"), getColor(999999, 999999) returns "x" immediately.

1. Add selectColumn(c). Rows and columns overlap without nesting, so "most specific wins" no longer works. A row paint and a column paint can both cover a cell, and only time decides. Switch to timestamps: keep a counter, store (stamp, color) for the table, each row, each column and each cell, and return the color with the largest stamp among the four. Clearing on a table paint still works. A row paint can still drop that row's cells, but it cannot drop column entries.
2. Select a rectangle selectRange(r1, c1, r2, c2). Store each range paint as (stamp, r1, c1, r2, c2, color). A simple getColor scans the paints from newest to oldest and returns the first one that contains the cell: O(P). If the interviewer wants faster queries, discuss a 2D segment tree with lazy "assign" tags, or rebuilding a compressed grid over the distinct range edges. Ask which operation is hot before you pick one.
3. Multi-select (Shift+click). Make the selection a list of shapes. setColor applies the same color to each shape with the same stamp, or with stamps in click order. Ask what happens when shapes overlap; with one color it does not matter.
4. Undo and redo. Before each setColor, push what it overwrites: the old table color plus the cleared row and cell maps for a table paint, the old row entry plus the dropped cell map for a row paint, or the old cell entry for a cell paint. Undo restores them. Because the old maps are moved, not copied, each push is still O(1). Redo is a second stack that is cleared on any new setColor.
5. Insert or delete rows and columns. Fixed integer indices break. Give each row and column a stable id and keep an ordered list from index to id. Store colors by id. Then inserting a row does not move any stored color.

Mid-level (L4)

• Asks what setColor does with no selection and how to treat bad indices.
• Writes a correct brute-force grid quickly, then moves to per-level storage when asked about large tables.
• Gets "last write wins" right, including a row paint after a cell paint.

Senior (L5)

• Starts from the cost of a whole-table paint and goes straight to O(1) operations.
• States the invariant (each stored entry is newer than the coarser ones above it) and explains why it makes the lookup order correct.
• Handles selectColumn by switching to timestamps, and explains why the nesting argument breaks.

Staff (L6+)

• Discusses memory growth and when to clear or compact entries, and why the clears are amortized O(1).
• Designs undo, range selection and multi-select without rewriting the core, and names the trade-offs between a paint log and a 2D segment tree.
• Connects the model to the real product: stable row and column ids, many users editing at once, and how last-write-wins maps to a shared document.