Generated by Codex with GPT 5.6 Sol High

Quick facts

• Difficulty: HARD
• Problem: Custom LeetCode-style prompt: Count Nodes in a Distributed Tree
• Topics: Distributed Systems, Tree, Asynchronous Programming, State Machine, Idempotency

Problem gist

Every tree node runs in a separate process or on a separate machine. A node knows only its parent and its direct children, and it can communicate with them only by sending asynchronous messages.

The root must eventually report the number of nodes in the whole tree. A leaf can answer 1 immediately. Every other node must:

1. ask each child to count its subtree
2. remember which children have replied
3. add the child results to its own count of 1
4. reply to its parent only after every child has answered

The arithmetic is easy. The real problem is coordinating asynchronous events without a call stack, a shared tree object, or blocking while waiting for a child.

The network follow-up makes the protocol more realistic: requests and responses may be duplicated or lost. The result must still be correct without counting a child twice.

Clarify the contract first

Before writing code, establish the model that the prompt leaves implicit:

• The topology is a rooted tree and does not change during one count operation.
• Every node has a stable unique ID and knows its parent and children.
• Sending a message returns immediately; a later event invokes the receiver.
• Processes do not crash or forget their state during the operation.
• A message may be lost, duplicated, delayed, or reordered.
• If a message is retried forever, some copy is eventually delivered.
• Several count operations may overlap, so every operation needs a unique ID.

The eventual-delivery assumption is essential for liveness. No algorithm can guarantee completion if every message to one child can be lost forever. The other assumptions can be changed, but each change requires more protocol machinery.

The reliable-network state machine

For one count request, a node stores:

• the request ID
• the parent that requested the count
• the set of children still pending
• one accepted count from each child
• the completed subtree count, once available

When a new request arrives, the node creates this state and sends a request to every child. It does not loop waiting for replies. Its message handler returns, and later response events advance the stored state.

When a child response arrives, the node records it and removes that child from the pending set. If the set becomes empty, the node computes:

subtree count = 1 + sum of all child subtree counts

It then sends that result to its parent. A leaf starts with an empty pending set, so it completes immediately with 1.

This is a distributed postorder traversal. Requests flow downward, and subtree counts flow upward. The ordinary recursive call stack has been replaced by persistent state at each node.

Making retries and duplicates safe

The root creates one globally unique request ID for the entire count operation. Every descendant forwards that same ID. Each node keys its state by this ID.

Duplicate requests are idempotent:

• if the request is still running, the node does not create another operation or reset partial progress
• if the request has finished, the node resends its cached result

Duplicate responses are also idempotent. A parent accepts a response only if that child is still in the pending set. Later responses from the same child and request ID are ignored.

Loss is handled with retries. While children remain pending, a node periodically resends requests only to those children. This repairs both directions:

• if the original request was lost, a retry eventually reaches the child
• if the child’s response was lost, the retry reaches a completed child, which replays its cached response

The protocol therefore provides effectively-once aggregation over an at-least-once retry process. The network may deliver a logical message many times, but each parent applies each child’s contribution once.

Why the result is correct

For a leaf, the stored result is 1, which is exactly the size of its subtree.

Assume every child of some node eventually reports its correct subtree size. The node accepts one result per child and computes 1 for itself plus the sum of those disjoint child subtrees. That is exactly the size of its subtree.

By induction from the leaves to the root, the root’s completed result is the number of nodes in the whole tree.

Retries do not change that proof. A duplicated child response cannot be added again after the child leaves the pending set, and a duplicated request cannot create a second state for the same request ID.

Python solution

from dataclasses import dataclass, field
from typing import Callable, Protocol, TypeAlias
from uuid import uuid4


NodeId: TypeAlias = str
RequestId: TypeAlias = str


@dataclass(frozen=True)
class CountRequest:
    request_id: RequestId


@dataclass(frozen=True)
class CountResponse:
    request_id: RequestId
    subtree_count: int


Message: TypeAlias = CountRequest | CountResponse


class AsyncTransport(Protocol):
    """Accept a message for asynchronous delivery and return immediately."""

    def send(
        self,
        sender_id: NodeId,
        recipient_id: NodeId,
        message: Message,
    ) -> None:
        ...


class RetryScheduler(Protocol):
    """Schedule a callback without blocking the current event handler."""

    def call_later(
        self,
        delay_seconds: float,
        callback: Callable[[], None],
    ) -> None:
        ...


@dataclass
class RequestState:
    requester_id: NodeId | None
    pending_children: set[NodeId]
    child_counts: dict[NodeId, int] = field(default_factory=dict)
    completed_count: int | None = None
    root_callback: Callable[[int], None] | None = None
    retry_scheduled: bool = False


class DistributedTreeNode:
    """
    Event-driven subtree counter for one node in a static distributed tree.

    The transport invokes receive() when a message is delivered. No method
    blocks while waiting for another process.
    """

    def __init__(
        self,
        node_id: NodeId,
        parent_id: NodeId | None,
        child_ids: set[NodeId],
        transport: AsyncTransport,
        scheduler: RetryScheduler,
        retry_delay_seconds: float = 1.0,
    ) -> None:
        if node_id in child_ids:
            raise ValueError("A node cannot be its own child")
        if retry_delay_seconds <= 0:
            raise ValueError("Retry delay must be positive")

        self.node_id = node_id
        self.parent_id = parent_id
        self.child_ids = frozenset(child_ids)
        self._transport = transport
        self._scheduler = scheduler
        self._retry_delay_seconds = retry_delay_seconds
        self._requests: dict[RequestId, RequestState] = {}

    def start_count(self, on_complete: Callable[[int], None]) -> RequestId:
        """Start a count at the root and return immediately."""
        if self.parent_id is not None:
            raise RuntimeError("Only the root can start a whole-tree count")

        request_id = str(uuid4())
        self._begin_request(
            request_id=request_id,
            requester_id=None,
            root_callback=on_complete,
        )
        return request_id

    def receive(self, sender_id: NodeId, message: Message) -> None:
        """Advance the state machine for one asynchronously delivered message."""
        if isinstance(message, CountRequest):
            self._handle_count_request(sender_id, message)
        elif isinstance(message, CountResponse):
            self._handle_count_response(sender_id, message)
        else:
            raise TypeError(f"Unsupported message type: {type(message).__name__}")

    def _handle_count_request(
        self,
        sender_id: NodeId,
        message: CountRequest,
    ) -> None:
        if sender_id != self.parent_id:
            raise ValueError("Count requests must come from this node's parent")

        state = self._requests.get(message.request_id)
        if state is None:
            self._begin_request(
                request_id=message.request_id,
                requester_id=sender_id,
                root_callback=None,
            )
            return

        if state.requester_id != sender_id:
            raise ValueError("A request ID cannot have two different requesters")

        if state.completed_count is not None:
            # The parent's earlier copy or our earlier response may have been
            # lost. Replaying a cached result makes the request idempotent.
            self._send_response(
                recipient_id=sender_id,
                request_id=message.request_id,
                subtree_count=state.completed_count,
            )
            return

        # A duplicate request can also accelerate recovery below this node.
        self._request_missing_children(message.request_id, state)
        self._ensure_retry_scheduled(message.request_id, state)

    def _handle_count_response(
        self,
        sender_id: NodeId,
        message: CountResponse,
    ) -> None:
        state = self._requests.get(message.request_id)

        # Ignore a stale response for an operation this process no longer knows.
        if state is None:
            return
        if sender_id not in self.child_ids:
            raise ValueError("Count responses must come from a direct child")
        if message.subtree_count < 1:
            raise ValueError("Every subtree count must be positive")

        # The child is removed after its first accepted response. Any later
        # duplicate therefore has no effect.
        if sender_id not in state.pending_children:
            return

        state.child_counts[sender_id] = message.subtree_count
        state.pending_children.remove(sender_id)

        if not state.pending_children:
            self._complete_request(message.request_id, state)

    def _begin_request(
        self,
        request_id: RequestId,
        requester_id: NodeId | None,
        root_callback: Callable[[int], None] | None,
    ) -> None:
        if request_id in self._requests:
            raise ValueError("Request already exists")

        state = RequestState(
            requester_id=requester_id,
            pending_children=set(self.child_ids),
            root_callback=root_callback,
        )
        self._requests[request_id] = state

        if not state.pending_children:
            self._complete_request(request_id, state)
            return

        self._request_missing_children(request_id, state)
        self._ensure_retry_scheduled(request_id, state)

    def _request_missing_children(
        self,
        request_id: RequestId,
        state: RequestState,
    ) -> None:
        # Iterate over a snapshot because a test transport may deliver messages
        # synchronously even though production delivery should be asynchronous.
        for child_id in sorted(state.pending_children):
            self._transport.send(
                self.node_id,
                child_id,
                CountRequest(request_id),
            )

    def _ensure_retry_scheduled(
        self,
        request_id: RequestId,
        state: RequestState,
    ) -> None:
        if (
            state.completed_count is not None
            or not state.pending_children
            or state.retry_scheduled
        ):
            return

        state.retry_scheduled = True
        self._scheduler.call_later(
            self._retry_delay_seconds,
            lambda: self._on_retry_timer(request_id),
        )

    def _on_retry_timer(self, request_id: RequestId) -> None:
        state = self._requests.get(request_id)
        if state is None:
            return

        state.retry_scheduled = False
        if state.completed_count is not None or not state.pending_children:
            return

        self._request_missing_children(request_id, state)
        self._ensure_retry_scheduled(request_id, state)

    def _complete_request(
        self,
        request_id: RequestId,
        state: RequestState,
    ) -> None:
        if state.completed_count is not None:
            return

        completed_count = 1 + sum(state.child_counts.values())

        # Cache completion before producing any external effect. Reentrant or
        # duplicated events can now only replay the same immutable result.
        state.completed_count = completed_count

        if state.requester_id is not None:
            self._send_response(
                recipient_id=state.requester_id,
                request_id=request_id,
                subtree_count=completed_count,
            )
        elif state.root_callback is not None:
            state.root_callback(completed_count)

    def _send_response(
        self,
        recipient_id: NodeId,
        request_id: RequestId,
        subtree_count: int,
    ) -> None:
        self._transport.send(
            self.node_id,
            recipient_id,
            CountResponse(request_id, subtree_count),
        )

AsyncTransport and RetryScheduler deliberately hide the runtime. They can be implemented with an actor framework, an event loop, RPC callbacks, a message broker, or a deterministic fake in tests. The node logic remains an ordinary single-threaded event-driven state machine.

On a reliable network, a tree with n nodes sends one request and one response across each of its n - 1 edges, for 2(n - 1) messages. The total local work is O(n), and the completion latency follows the tree height because sibling subtrees run concurrently.

For one active request, a node stores O(number of children) state. Across the tree that is O(n). With loss and duplication, there is no fixed message bound; the cost is O(n + r), where r is the number of retry and duplicate deliveries before completion.

A small message-flow example

Suppose root A has children B and C, and B has leaf child D.

A -> B: CountRequest(R)
A -> C: CountRequest(R)
B -> D: CountRequest(R)
C -> A: CountResponse(R, 1)
D -> B: CountResponse(R, 1)
B -> A: CountResponse(R, 2)
A completes with 1 + 1 + 2 = 4

The responses may arrive in any order. If B’s response is lost, A eventually resends CountRequest(R) to B. Because B cached 2 for request R, it immediately sends the same response again instead of recounting D.

Interview follow-ups

What if messages can be duplicated?

Attach the root-generated request ID to every message and store request state by that ID. A node must not restart work when it sees the same request again. A parent also keys accepted responses by child ID, represented in the code by the pending-child set.

This makes every state transition idempotent. The first response from a child changes the aggregate; later copies do nothing. The first request creates the state; later copies either nudge pending retries or replay the completed result.

Using only a request ID is not enough if the parent adds every response it sees. Deduplication needs the pair (request_id, child_id).

What if messages can be lost?

Use timers and retry only the children still missing for that request. A retry must reuse the same request ID; generating a new one would defeat deduplication and could start multiple independent counts.

The parent-driven retry in this solution repairs lost responses as well as lost requests. If the response disappeared, the repeated request causes the child to replay its cached result. Exponential backoff with jitter is preferable to a fixed interval in a real deployment because it reduces synchronized retry storms.

Retries provide liveness only under an eventual-delivery assumption. If a child or network partition can remain unavailable forever, the API also needs a deadline and a defined failure result.

Can this protocol claim exactly-once message delivery?

No. In an unreliable network, the sender cannot distinguish a lost message from a delivered message whose reply was lost. Retrying can therefore create duplicates.

The useful guarantee is effectively-once processing: delivery is at least once eventually, while idempotent handlers and deduplication ensure each child’s logical contribution affects the total once. This distinction is important in distributed-system interviews.

How can completed request state be garbage-collected?

The simple solution retains completed results so an arbitrarily late retry can be answered safely. That is correct but unbounded.

A production protocol can add a CountAck(request_id) message. A child retains and retries its completed response until its parent acknowledges it. The parent must acknowledge duplicate responses too. After receiving the acknowledgment, the child can delete the cached result, possibly after an additional retention window for very delayed duplicates.

Another option is a time-to-live longer than the maximum retry and message lifetime. That is simpler but makes correctness depend on a timing bound. A durable deduplication store may be needed when request IDs must remain safe across process restarts.

What if a node crashes and restarts?

The in-memory implementation loses partial state, so process failure is outside its guarantee. To tolerate restarts, persist the request state and completed result before acknowledging or sending transitions that depend on them. A write-ahead log, transactional local database, or replicated actor state can provide this durability.

If a restarted node receives a request whose result was durably completed, it replays the result. If the operation was incomplete, it restores the pending children and resumes retries. The same idempotent protocol still applies; only the state storage changes.

What if the tree changes while counting?

Without a consistency rule, the answer may mix two different topologies: a node could join after its parent’s child list was captured, or move while an old request is still running.

The cleanest interview assumption is that topology is fixed for one operation. If mutations must continue, tag the count with a topology version or run it against a consistent snapshot. More advanced systems may use barriers, distributed snapshots, or membership epochs. The desired semantics must be defined before choosing among them.

What if several count requests run at the same time?

The state map already isolates them by request ID. Each request has its own pending set, partial counts, retries, and cached result.

This costs O(number of concurrent requests * local degree) memory per node. Production code should enforce concurrency limits, reject duplicate client operations with an idempotency key, and expire abandoned requests so a caller cannot exhaust memory.

How would timeouts or partial results work?

A deadline can terminate retries and report which child subtrees are missing. The result should not be presented as an exact node count. A useful response could contain the confirmed count, the missing child IDs, and an explicit complete = false flag.

Cancellation is another distributed message and can also be lost or duplicated, so it should carry the request ID and be idempotent. Nodes can release active state after cancellation while retaining a short-lived tombstone to prevent a late request from accidentally restarting the operation.

Could the messages be combined to reduce traffic?

For one exact count, the reliable case is already asymptotically optimal: every non-root node must influence the root, so information has to cross every tree edge. Aggregating at each node reduces upward traffic to one integer response per edge instead of sending every descendant identity to the root.

If counts are requested frequently and the tree changes slowly, each node can maintain a cached subtree count and update ancestors incrementally when members join or leave. That improves read latency but shifts complexity to keeping the cache consistent during concurrent topology changes and failures.

Practical takeaway

The recursive equation is still 1 + sum(child counts), but recursion is not the implementation model. In a distributed setting, first identify message boundaries, durable per-request state, completion conditions, and failure semantics. Once those are explicit, the familiar tree algorithm becomes a small, testable asynchronous state machine.