Count and map a cluster tree with message passing

No node can see the whole tree. Each one knows only its parent and its children, and the only way to learn anything is to send a message and wait for the answer. So both operations work the same way: a scatter-gather (also called convergecast) down and back up the tree.

1. A request reaches a node. The node records who asked and which children it is waiting on, then forwards the same request to every child.
2. A leaf has nobody to wait for, so it answers right away.
3. A child's answer is folded into the node's partial result. When the last child has answered, the node sends its combined result to whoever asked it: its parent, or the client if it is the root.

Count. A node's answer is 1 + sum(child answers).

Topology. A node's answer is id(child_1,child_2,...), or just id for a leaf. Because answers can arrive in any order, keep each child's string in a map and join them in the node's own declared child order when finishing, not in arrival order. Otherwise the output would change from run to run.

Why a call stack doesn't work. The tempting first version is a recursive count() that calls the children and returns the sum. That assumes a synchronous RPC. Here send only enqueues a message and returns at once, and the answer comes back later through receive. All state must therefore live on the node between messages. That is the real point of the question, and it is where most candidates trip.

Invariants

• Every request carries a request id (rid). Per-node state is keyed by it, so a count and a topology (or two counts) can run at the same time without mixing.
• waiting holds exactly the children that haven't answered yet. The node sends its reply once, at the moment waiting becomes empty.
• A reply from a child that isn't in waiting (a duplicate or a late reply) is ignored, so a single child is never counted twice.
• Each tree edge carries exactly one request down and one reply up per operation.

from __future__ import annotations

import random
from dataclasses import dataclass, field
from typing import Callable, Optional

CLIENT = "-1"  # the root's "parent" is the external client


@dataclass
class Pending:
    """Per-request state while a node waits for its children."""
    kind: str
    reply_to: str
    waiting: set[str]
    count: int = 1                                   # this node itself
    parts: dict[str, str] = field(default_factory=dict)


class Network:
    """Provided transport: send() only enqueues; delivery order is arbitrary."""

    def __init__(self, seed: int = 0) -> None:
        self.nodes: dict[str, Node] = {}
        self.queue: list[tuple[str, str, dict]] = []
        self.rng = random.Random(seed)
        self.client_inbox: list[dict] = []

    def send(self, src: str, dst: str, msg: dict) -> None:
        self.queue.append((src, dst, msg))

    def run(self) -> None:
        while self.queue:
            src, dst, msg = self.queue.pop(self.rng.randrange(len(self.queue)))
            if dst == CLIENT:
                self.client_inbox.append(msg)
            else:
                self.nodes[dst].receive(src, msg)


class Node:
    def __init__(self, node_id: str, parent: str, children: list[str], net: Network) -> None:
        self.id = node_id
        self.parent = parent
        self.children = children
        self.net = net
        self.pending: dict[str, Pending] = {}        # request_id -> state
        net.nodes[node_id] = self

    # The only entry point: every message arrives here.
    def receive(self, src: str, msg: dict) -> None:
        kind, rid = msg["type"], msg["rid"]
        if kind in ("count", "topology"):
            self._start(kind, rid, reply_to=src)
        elif kind in ("count_result", "topology_result"):
            self._absorb(src, rid, msg["value"])
        else:
            raise ValueError(f"unknown message type {kind!r}")

    def _start(self, kind: str, rid: str, reply_to: str) -> None:
        if rid in self.pending:
            return                                   # duplicate request
        state = Pending(kind, reply_to, waiting=set(self.children))
        self.pending[rid] = state
        for child in self.children:
            self.net.send(self.id, child, {"type": kind, "rid": rid})
        if not self.children:                        # leaf answers at once
            self._finish(rid)

    def _absorb(self, child: str, rid: str, value: int | str) -> None:
        state = self.pending.get(rid)
        if state is None or child not in state.waiting:
            return                                   # duplicate or stale reply
        state.waiting.discard(child)
        if isinstance(value, int):
            state.count += value
        else:
            state.parts[child] = value
        if not state.waiting:
            self._finish(rid)

    def _finish(self, rid: str) -> None:
        state = self.pending.pop(rid)
        value: int | str
        if state.kind == "count":
            value = state.count
        else:
            inner = ",".join(state.parts[c] for c in self.children)  # declared order
            value = f"{self.id}({inner})" if inner else self.id
        self.net.send(self.id, state.reply_to, {"type": state.kind + "_result", "rid": rid, "value": value})


def build(parents: list[int], seed: int = 0) -> Network:
    """parents[i] is the parent of node i (-1 for the root)."""
    net = Network(seed)
    kids: dict[int, list[str]] = {i: [] for i in range(len(parents))}
    for i, p in enumerate(parents):
        if p >= 0:
            kids[p].append(str(i))
    for i, p in enumerate(parents):
        Node(str(i), str(p), kids[i], net)
    return net


def ask(net: Network, root: str, kind: str, rid: str = "r1") -> int | str:
    net.send(CLIENT, root, {"type": kind, "rid": rid})
    net.run()
    return next(m["value"] for m in net.client_inbox if m["rid"] == rid)

Network stands in for the provided messaging API. It delivers queued messages in random order, which shows the node logic doesn't depend on delivery order. In the interview you implement Node.receive and whatever state it needs; send is given.

For n machines and tree height h:

• Messages: 2(n − 1) per operation, one request and one reply on each edge.
• Latency: 2h message hops when messages travel in parallel.
• Count: O(1) work per message. A node holds O(children) state per in-flight request.
• Topology: the string at a node is O(size of its subtree). That makes the total bytes sent O(n · h) (O(n²) for a path), and the root holds O(n). The total work is O(n · h) for the same reason, since each level copies its children's strings.
• Space across the cluster: O(n) per in-flight request (each node keeps its waiting set and child parts until it finishes).

• Root 0 with children 1, 2; 1 has 3, 4; 2 has 5, 6 → count 7, topology 0(1(3,4),2(5,6)), for every delivery order (the tests try 20 random seeds).
• Single node → count 1, topology 0.
• A path 0 → 1 → 2 → 3 → count 4, topology 0(1(2(3))).
• The root isn't node 0: parents [2, 2, -1, 0] → count 4, topology 2(0(3),1).
• A count and a topology sent to the root at once with ids a and b → a: 4, b: 0(1(3),2), with no cross-talk.
• A star with 1,000 leaves → count 1001. The root sends 1,000 messages and folds in 1,000 replies.
• A duplicate reply from the same child → ignored, so the count is still correct.

1. Lost messages or dead machines. Add a timeout per request at each node. When it fires, either resend to the silent children (requests are idempotent by rid, and duplicate replies are already dropped) or give up and reply with a partial result plus the ids of the subtrees that didn't answer, for example {"count": 5, "missing": ["2"]}. Say clearly which behaviour you choose: an exact answer that can fail, or a best-effort answer that is always delivered.
2. Very deep trees (millions of machines). The topology string grows with the subtree, so messages near the root get large. Options: return (parent, child) edge lists streamed in chunks, send each subtree to a blob store and pass a reference up, or have each node report its edges straight to a collector. Count stays O(1) bytes per message.
3. Other aggregates. Sum of a per-node metric, max, the height of the tree, or the list of nodes with low disk space. Any associative merge fits in the same skeleton; only the leaf value and the merge function change. Generalize with an Aggregator interface (leaf(), merge(acc, child_value), finish(acc)).
4. Request from any node, not just the root. Treat the tree as an undirected graph. Forward the request to every neighbour except the one it came from (parent included), and the same convergecast still works. Alternatively, route the request up to the root first.
5. Topology changes while a request runs. A child that joins mid-request isn't in waiting, so it is simply not counted, which gives a consistent snapshot from the moment the request arrived. A child that leaves needs the timeout from follow-up 1. If exact answers are required, use versioned membership or a snapshot protocol.

Mid-level (L4)

• Sees that send is asynchronous and that results come back through receive, so state lives on the node rather than on a call stack.
• Gets count and topology right on a static tree, including leaves and a single-node cluster.
• Gets the parentheses and commas right in the topology format, and builds it deterministically.

Senior (L5)

• Adds request ids so that concurrent requests don't corrupt each other, and frees per-request state once the node finishes.
• Joins topology parts in declared child order rather than arrival order, and explains why.
• Counts messages and latency (2(n − 1) messages, 2h hops) and notices that topology payloads grow toward the root.
• Discusses lost messages and duplicate replies, and makes requests idempotent.

Staff (L6+)

• Separates the convergecast skeleton from the aggregate (an Aggregator interface), so new metrics need no protocol change.
• Makes the trade-offs explicit: exact or best-effort under failures, timeouts that grow with subtree depth, and snapshot semantics while membership changes.
• Fixes payload growth at scale (streamed edge lists, references to a blob store) and considers backpressure at high-fan-out nodes.