Code Editor with Block Shrink and Expand (Code Folding)
Company: Jane Street
Role: Machine Learning Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
Implement the core of a code editor (similar to VS Code) that supports **code folding**: collapsing ("shrink") and restoring ("expand") indentation-based blocks of code.
## Setup
The editor is initialized with a list of source-code lines. Indentation determines block structure, as in Python:
- Indentation is measured by the number of leading spaces, and every line's indentation is a multiple of 4 spaces.
- A line `L` is a **block header** if the line immediately after it has strictly greater indentation than `L`.
- The **block** owned by header `L` consists of all consecutive lines after `L` whose indentation is strictly greater than `L`'s indentation. The block ends at the first subsequent line whose indentation is less than or equal to `L`'s (or at the end of the file).
- Blocks may be nested: a line inside a block may itself be a header of a smaller block.
## Rendering
Implement a `render()` method that returns the currently **visible** lines, one string per line, formatted as:
```
<display_number> <marker><line content>
```
where:
- `display_number` numbers the visible lines sequentially starting from 1 (hidden lines are skipped and do not consume a number).
- `marker` is `+ ` if the line is a block header, and the empty string otherwise.
- `line content` is the original text of the line, including its leading indentation.
## Operations
Implement two operations. Both take a **display line number** — a number as shown in the most recent `render()` output, not the original file line number.
- `shrink(d)`: if the visible line currently displayed at number `d` is a block header, mark it **collapsed**. A collapsed header stays visible, but every line inside its block becomes hidden. If the line at `d` is not a block header, or its block is already collapsed, this is a no-op.
- `expand(d)`: if the visible line currently displayed at number `d` is a collapsed block header, mark it expanded again. Lines inside its block become visible **except** lines that are hidden because some nested header inside the block is itself still collapsed (each header keeps its own collapsed/expanded state).
In general, a line is visible if and only if none of the headers whose blocks contain it are collapsed.
## Example
The editor is initialized with these 5 lines:
```
a = 1
for i in range(10):
print(i)
print(i ** 2)
b = 10
```
`render()` returns:
```
1 a = 1
2 + for i in range(10):
3 print(i)
4 print(i ** 2)
5 b = 10
```
After calling `shrink(2)`, `render()` returns:
```
1 a = 1
2 + for i in range(10):
3 b = 10
```
Lines 3 and 4 of the original file (the body of the `for` loop) are hidden, and `b = 10` is renumbered to display number 3. Calling `expand(2)` afterwards restores the original output.
## Task
Implement a class:
```python
class CodeEditor:
def __init__(self, lines: list[str]): ...
def shrink(self, d: int) -> None: ...
def expand(self, d: int) -> None: ...
def render(self) -> list[str]: ...
```
## Constraints
- `1 <= number of lines <= 10^4`, each line at most 200 characters.
- Indentation is spaces only, always a multiple of 4, and the first line has indentation 0. A line's indentation is at most one level (4 spaces) deeper than the line before it.
- Up to `10^4` total `shrink` / `expand` / `render` calls.
- Operations always receive a display number `d` that is valid for the most recent rendering (i.e., `1 <= d <=` the current number of visible lines).
Quick Answer: This question evaluates a candidate's ability to model and manipulate indentation-based hierarchical structures, maintain nested visibility state for collapsed blocks, and map between original file lines and dynamic display numbering.
Implement the core of a code editor (like VS Code) that supports **code folding**: collapsing ("shrink") and restoring ("expand") indentation-based blocks of code.
Indentation (leading spaces, always a multiple of 4) defines block structure, as in Python. Line `i` is a **block header** if the line immediately after it has strictly greater indentation. The block owned by a header is the run of consecutive following lines with strictly greater indentation (it ends at the first line whose indentation is <= the header's). Blocks may nest, and each header keeps its own collapsed/expanded state. A line is visible if and only if none of the headers whose blocks contain it are collapsed.
### `render()` format
`render()` returns the currently visible lines, one string each, as:
```
<display_number> <marker><line content>
```
- `display_number`: 1-based index over visible lines only (hidden lines consume no number).
- `marker`: `"+ "` if the line is a block header, otherwise the empty string.
- `line content`: the original text, including leading indentation.
### Operations (display numbers, not file line numbers)
- `shrink(d)`: if the visible line at display number `d` is a block header, mark it collapsed (its block's lines become hidden; the header stays visible). No-op if it is not a header or is already collapsed.
- `expand(d)`: if the visible line at `d` is a collapsed header, expand it. Its block's lines become visible **except** those still hidden by a nested header that remains collapsed.
### Function adaptation (what you implement here)
This stateful editor is exposed as a single function so it can be graded by replay:
```
solution(lines, operations) -> list of render outputs
```
- `lines`: the source lines the editor is initialized with.
- `operations`: a sequence of operations applied in order, each one of:
- `["render"]` — append the current `render()` output to the result,
- `["shrink", d]` — collapse the header at display number `d`,
- `["expand", d]` — expand the collapsed header at display number `d`.
- Return the list of `render()` outputs (a list of lists of strings), one entry per `["render"]` operation, in the order they occur. A `shrink`/`expand` uses the display numbering of the current (most recent) render state.
### Example
Initialized with:
```
a = 1
for i in range(10):
print(i)
print(i ** 2)
b = 10
```
`render()` returns:
```
1 a = 1
2 + for i in range(10):
3 print(i)
4 print(i ** 2)
5 b = 10
```
After `shrink(2)`, `render()` returns:
```
1 a = 1
2 + for i in range(10):
3 b = 10
```
The for-loop body (original lines 3-4) is hidden and `b = 10` is renumbered to 3. `expand(2)` restores the original output.
Constraints
- 1 <= number of lines <= 10^4; each line has at most 200 characters.
- Indentation is spaces only, always a multiple of 4; the first line has indentation 0.
- A line's indentation is at most one level (4 spaces) deeper than the previous line.
- Up to 10^4 total shrink / expand / render operations.
- Every shrink/expand receives a display number d that is valid for the current render state (1 <= d <= number of visible lines).
Examples
Input: (['a = 1', 'for i in range(10):', ' print(i)', ' print(i ** 2)', 'b = 10'], [['render'], ['shrink', 2], ['render'], ['expand', 2], ['render']])
Expected Output: [['1 a = 1', '2 + for i in range(10):', '3 print(i)', '4 print(i ** 2)', '5 b = 10'], ['1 a = 1', '2 + for i in range(10):', '3 b = 10'], ['1 a = 1', '2 + for i in range(10):', '3 print(i)', '4 print(i ** 2)', '5 b = 10']]
Explanation: Canonical example: render, shrink the for-loop, render (body hidden + renumbered), expand, render (restored).
Input: (['def f():', ' for i in x:', ' a = 1', ' b = 2', ' c = 3', 'd = 4'], [['render'], ['shrink', 2], ['render'], ['shrink', 1], ['render'], ['expand', 1], ['render']])
Expected Output: [['1 + def f():', '2 + for i in x:', '3 a = 1', '4 b = 2', '5 c = 3', '6 d = 4'], ['1 + def f():', '2 + for i in x:', '3 c = 3', '4 d = 4'], ['1 + def f():', '2 d = 4'], ['1 + def f():', '2 + for i in x:', '3 c = 3', '4 d = 4']]
Explanation: Nested blocks: collapse inner for-loop, collapse outer def, expand outer while inner stays collapsed.
Hints
- Precompute each line's indentation (leading spaces) and whether it is a block header: line i is a header iff line i+1 has strictly greater indentation.
- Store a boolean collapsed flag per line. A line is visible iff none of the headers whose blocks contain it are collapsed.
- To map a display number to a real line and to render, scan top-to-bottom: when you reach a visible collapsed header, jump past every following line whose indentation is greater than the header's (its entire block, including any nested headers).
- shrink/expand only toggle the targeted header's own flag. Nested headers keep independent state, which is why expanding an outer block leaves inner collapsed blocks folded.