Customer revenue with referral credit and bottom-k query

Each customer has two numbers: own revenue (what they generated) and total revenue (own plus the own revenue of the customers they referred directly). Referral credit goes up exactly one level. If C was referred by B and B by A, C's revenue counts toward B's total but not A's.

Brute force. Store own and referrer for each customer. For a query, rebuild every total in O(n), filter by the threshold, sort and take k. That is O(n log n) per query, which is fine for a handful of queries. It is too slow when queries and inserts interleave at scale, and interviewers usually ask about that next.

Optimal. Keep the totals up to date as revenue comes in, and keep a sorted index of (total, id) pairs:

1. add_customer(revenue, referrer_id) gives out the next integer id, records own and total revenue, and inserts (revenue, id) into the index. If there is a referrer, its total grows by revenue. We remove its old (total, id) entry and insert the new one.
2. add_revenue(id, amount) (a later sale) raises the customer's own and total revenue and the referrer's total by amount. That re-indexes at most two entries.
3. lowest_k(k, min_revenue) binary-searches the first entry whose total is strictly greater than min_revenue and returns the next k ids. Ties are broken by id because the index is ordered by (total, id).

Invariants:

• total[c] == own[c] + sum(own[r] for r in children[c]) for every customer after every operation.
• ranked holds exactly one (total[c], c) per customer and stays sorted. Every total change is a remove followed by an insert, never an in-place edit.
• A single event changes at most two totals (the customer and their direct referrer). That bound is what keeps updates cheap.

In Python, bisect plus insort on a list is the simplest version. The search is O(log n), but list inserts and deletes shift memory, so they cost O(n) in the worst case (a fast memmove in practice). For true O(log n), say you would use sortedcontainers.SortedList or a balanced BST / skip list (TreeSet in Java, std::set in C++). The interface does not change.

from bisect import bisect_right, insort
from collections import deque
from typing import Optional


class RevenueSystem:
    """Tracks revenue per customer; a referrer is credited with the
    revenue of every customer they directly referred."""

    def __init__(self) -> None:
        self._own: list[int] = []                # revenue the customer generated
        self._total: list[int] = []              # own + direct referrals' own
        self._referrer: list[Optional[int]] = []
        self._children: list[list[int]] = []
        self._ranked: list[tuple[int, int]] = [] # sorted (total, id)

    def _set_total(self, cid: int, new_total: int) -> None:
        old = (self._total[cid], cid)
        i = bisect_right(self._ranked, old) - 1  # exact entry is present
        del self._ranked[i]
        self._total[cid] = new_total
        insort(self._ranked, (new_total, cid))

    def add_customer(self, revenue: int, referrer_id: Optional[int] = None) -> int:
        if revenue < 0:
            raise ValueError("revenue must be non-negative")
        if referrer_id is not None and not 0 <= referrer_id < len(self._own):
            raise KeyError(f"unknown referrer {referrer_id}")
        cid = len(self._own)
        self._own.append(revenue)
        self._total.append(revenue)
        self._referrer.append(referrer_id)
        self._children.append([])
        insort(self._ranked, (revenue, cid))
        if referrer_id is not None:
            self._children[referrer_id].append(cid)
            self._set_total(referrer_id, self._total[referrer_id] + revenue)
        return cid

    def add_revenue(self, customer_id: int, amount: int) -> None:
        if not 0 <= customer_id < len(self._own):
            raise KeyError(f"unknown customer {customer_id}")
        if amount < 0:
            raise ValueError("amount must be non-negative")
        self._own[customer_id] += amount
        self._set_total(customer_id, self._total[customer_id] + amount)
        ref = self._referrer[customer_id]
        if ref is not None:
            self._set_total(ref, self._total[ref] + amount)

    def total_revenue(self, customer_id: int) -> int:
        return self._total[customer_id]

    def lowest_k(self, k: int, min_revenue: int) -> list[int]:
        """Ids of the k customers with the smallest total revenue strictly
        greater than min_revenue, ascending by (total, id)."""
        if k <= 0:
            return []
        start = bisect_right(self._ranked, (min_revenue, float("inf")))
        return [cid for _, cid in self._ranked[start:start + k]]

    def nested_revenue(self, customer_id: int, max_depth: int) -> int:
        """Own revenue plus revenue of everyone up to max_depth referral
        levels below (depth 1 == total_revenue)."""
        total, frontier = 0, deque([(customer_id, 0)])
        while frontier:
            cid, depth = frontier.popleft()
            total += self._own[cid]
            if depth < max_depth:
                frontier.extend((c, depth + 1) for c in self._children[cid])
        return total

Let n be the number of customers. Costs with the plain sorted list are shown first, and costs with a SortedList / balanced BST in brackets.

• add_customer: O(n) [O(log n)]. One insert, plus one remove and one insert for the referrer.
• add_revenue: O(n) [O(log n)]. At most two re-indexes.
• total_revenue: O(1).
• lowest_k: O(log n + k). One binary search plus a slice of k.
• nested_revenue(id, d): O(size of the referral subtree within depth d).
• Space: O(n) for the per-customer arrays, the child lists and the index.

• a = add_customer(10), b = add_customer(20, a), c = add_customer(5), d = add_customer(15, b): totals are a=30, b=35, c=5, d=15. A's total does not include D, who is two levels down.
• lowest_k(2, 10) gives [d, a] (15 and 30; c=5 is filtered out).
• lowest_k(10, 0) gives [c, d, a, b]. k larger than the number of matches returns all of them.
• lowest_k(3, 35) gives []. The threshold is strict, so b's 35 is excluded.
• lowest_k(0, 0) gives [].
• add_revenue(d, 30): d becomes 45 and b becomes 65. a stays 30 because credit only goes up one level. lowest_k(2, 5) then gives [a, d].
• nested_revenue(a, 0) = 10, nested_revenue(a, 1) = 30 (the same as the total), nested_revenue(a, 2) = 75.
• Ties: two customers with 7 each, lowest_k(5, 6) gives [x, y] in id order.
• An unknown referrer id raises KeyError. Negative revenue raises ValueError.

• Nested (indirect) revenue to depth d. nested_revenue walks the referral tree with BFS, which works for occasional queries. If you need to rank by nested revenue, keep a total for each depth up to d. Each event then walks up at most d ancestors and re-indexes each one, so it costs O(d log n). With unlimited depth this becomes subtree sums. Use an Euler tour plus a Fenwick tree if referrers are known up front, or walk the ancestor chain and accept O(depth) per event.
• Top k instead of bottom k, or ranges. Use the same index and binary-search from the other end. "Every customer with total in [lo, hi]" is two bisects and a slice.
• Revenue arrives as a stream of events, possibly with refunds. add_revenue already handles later sales. Negative amounts work once you drop the non-negative check. Re-indexing is unchanged because we always remove the old key before inserting the new one.
• Referrer changes or a customer is deleted. Subtract the customer's own revenue from the old referrer and add it to the new one. On delete, also remove the customer's index entry and decide what happens to their referrals (orphan them or re-parent them).
• Scale and concurrency. Totals are sums, so shard events by customer and aggregate. Bottom-k becomes a per-shard bottom-k followed by a merge. For a single process, put a lock around each event so the two-total update is atomic with respect to queries.

Mid-level (L4)

• Defines total revenue correctly (one level only) and handles a customer with no referrer.
• Gets a working version, even the rebuild-and-sort one, then improves updates to O(1) totals.
• Handles the strict threshold, k larger than the number of matches, and ties.

Senior (L5)

• Keeps a sorted index up to date with remove-then-insert, and states its invariants.
• Knows the real cost of insort on a Python list and names the O(log n) structure they would use.
• Adds add_revenue and nested revenue cleanly without rewriting the core.

Staff (L6+)

• Discusses the trade-off between maintaining totals for each depth and computing nested revenue on demand, and the subtree-sum approach for unlimited depth.
• Covers refunds, referrer changes and deletes without breaking the invariants.
• Sketches a sharded or streaming design (per-shard bottom-k plus a merge) and atomicity for the two-total update.