Ai2 Senior Software Engineer Interview Experience — A Streaming Step-Hierarchy Design Problem

The question gives a set of StepProgressEvents, where each event contains:

timestamp: datetime
step_id: int
parent_step_id: Optional[int]
run_state: RunState   # running / done / ...
text: str

The same step_id can receive multiple updates — for example first running, then later done — and events can also arrive out of order. Each step can have child steps, and in the end you need to maintain and output the current latest step hierarchy, for example:

1 done - Query intent: literature review
2 running - Retrieving candidate papers
3 done - Found 21 candidate papers by keyword search
4 running - Semantic search
5 running - Embedding query terms

Also, after the initial batch of events is processed, the system still needs to keep receiving arbitrary new events and update the tree in real time.

My solution:

• Define a StepNode: stores that step's latest_event and children.
• StepTracker maintains:
  • nodes: step_id -> StepNode, for O(1) lookup and update;
  • roots, holding the nodes where parent_step_id is None;
  • pending_children: parent_id -> children, to handle the case where a child arrives before its parent.
• All initial events also go through add_event(event) uniformly, so batch and streaming don't need two separate code paths.
• Inside add_event:
  1. Find or create the node by step_id;
  2. If the event's timestamp is older than the currently stored version, ignore it;
  3. Otherwise update the node's latest state with this event;
  4. If the parent already exists, attach directly under the parent's children; if the parent hasn't arrived yet, put it in pending_children;
  5. If, once the current node arrives, there are children that were waiting for it, attach them.
• Finally, starting from roots, use an iterative DFS with a stack to output the hierarchy. The stack holds (node, depth), using depth to control indentation and avoid recursion-depth issues.

Complexity:

• Updating a single event: O(1) on average
• Outputting the whole current tree: O(N)

I think what this question is really testing isn't the DFS itself, but:

1. Under out-of-order events, you must keep the latest state for the same step_id based on timestamp;
2. The arrival order of parent / child is not guaranteed;
3. Designing it as an incremental update, rather than rebuilding the whole tree from scratch on every new event.