Referral credit tracker with ranking queries: Frontier Interview's solution

Question: Referral credit tracker with ranking queries

Companies: Databricks, Uber

Difficulty: Medium

Problem statement

Build an in-memory tracker for account credit where new accounts can name a referrer.

Accounts get integer IDs in the order they are created, starting at 0. Each account has a total credit: its own starting credit, plus the starting credit of every account that named it as its direct referrer. Credit goes up one level only. If C names B and B named A, C's credit goes to B, not to A.

Implement AccountCreditTracker:

AccountCreditTracker() creates an empty tracker.
add_account(credit: float) -> int creates an account with that starting credit and returns its ID.
add_account_with_referral(credit: float, referrer_id: int) -> int creates an account with that starting credit, adds the same amount to the total of account referrer_id, and returns the new ID. The referrer must already exist.
get_total_credit(account_id: int) -> float returns the account's total credit.
top_k(k: int) -> list[int] returns the IDs of the k accounts with the highest total credit, highest first. On a tie, the lower ID comes first. If there are fewer than k accounts, return all of them.
count_at_least(threshold: float) -> int returns how many accounts have a total credit of at least threshold.

Constraints: credit amounts are non-negative. There can be up to 10^5 operations, mixed freely between writes and queries, so queries should not re-sort every account.

Example

t = AccountCreditTracker()
t.add_account(20)                    # -> 0
t.add_account(35)                    # -> 1
t.add_account_with_referral(25, 0)   # -> 2; account 0 total = 45
t.add_account_with_referral(5, 2)    # -> 3; account 2 total = 30, account 0 stays 45
t.top_k(3)                           # -> [0, 1, 2]   (45, 35, 30)
t.count_at_least(30)                 # -> 3
t.get_total_credit(3)                # -> 5

Expect follow-ups such as credit that flows up the whole referral chain, or credit that changes after an account is created.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Every account has a total credit: its own starting credit plus the
starting credit of each account that named it as the direct referrer.
A referral is one level only. If C was referred by B and B by A, C's credit
goes to B and not to A.

IDs come out in creation order, so plain arrays indexed by ID are enough
for the per-account state. No hash maps are needed. A referrer must already
exist, so the referral graph is a forest whose edges always point to a
smaller ID.

Brute force. Store each account's starting credit and referrer. To read
a total, scan all accounts and add up the ones that point at it. That is
O(n) per read and O(n log n) for a ranking query, which is fine for a first
pass but slow once reads dominate.

Incremental totals. When an account is created with a referrer, add its
credit to the referrer's total right away. Creation is O(1) and a total read
is O(1). This is all the first part asks for.

Ranking queries. The usual next part asks for the top k accounts by
total credit, or how many accounts are above a threshold. Recomputing with
heapq.nlargest costs O(n log k) per query. Keeping a sorted index of
(-total, id) pairs instead makes top_k a slice and count_at_least one
binary search. Each write moves at most two entries in that index: the new
account and its referrer. Negating the total puts the largest first, and the
ID breaks ties toward the older account. Both behaviours are deterministic,
so state them out loud.

Invariants:

total[i] = credit[i] + sum(credit[j] for j whose referrer is i).
_ranked holds exactly one (-total[i], i) per account, in sorted order.
An account's referrer has a smaller ID than the account.

Code

from __future__ import annotations

from bisect import bisect_right, insort


class AccountCreditTracker:
    """Accounts get IDs 0, 1, 2, ... in creation order.

    total[i] = account i's own initial credit + the initial credit of every
    account that named i as its direct referrer.
    """

    def __init__(self) -> None:
        self._total: list[float] = []        # own + direct referrals
        self._referrer: list[int | None] = []
        self._referrals: list[int] = []      # count of direct referrals
        # Sorted by (-total, id): highest credit first, ties by lower ID.
        self._ranked: list[tuple[float, int]] = []

    # ---- writes -------------------------------------------------------
    def add_account(self, credit: float) -> int:
        return self._create(credit, None)

    def add_account_with_referral(self, credit: float, referrer_id: int) -> int:
        if not 0 <= referrer_id < len(self._total):
            raise ValueError(f"unknown referrer {referrer_id}")
        return self._create(credit, referrer_id)

    def _create(self, credit: float, referrer_id: int | None) -> int:
        if credit < 0:
            raise ValueError("credit must be non-negative")
        new_id = len(self._total)
        self._total.append(credit)
        self._referrer.append(referrer_id)
        self._referrals.append(0)
        insort(self._ranked, (-credit, new_id))
        if referrer_id is not None:
            self._referrals[referrer_id] += 1
            self._bump(referrer_id, credit)
        return new_id

    def _bump(self, account_id: int, delta: float) -> None:
        old = (-self._total[account_id], account_id)
        pos = bisect_right(self._ranked, old) - 1   # exact key is present
        self._ranked.pop(pos)
        self._total[account_id] += delta
        insort(self._ranked, (-self._total[account_id], account_id))

    # ---- reads --------------------------------------------------------
    def get_total_credit(self, account_id: int) -> float:
        return self._total[account_id]

    def top_k(self, k: int) -> list[int]:
        """IDs of the k accounts with the most total credit, ties by lower ID."""
        return [account_id for _, account_id in self._ranked[:max(k, 0)]]

    def count_at_least(self, threshold: float) -> int:
        """How many accounts have total credit >= threshold."""
        # Entries with -total <= -threshold, i.e. total >= threshold.
        return bisect_right(self._ranked, (-threshold, float("inf")))

_bump finds the old key with bisect_right(...) - 1. The key is unique
because it contains the ID, so this always lands on that exact entry.

Complexity

Let n be the number of accounts.

add_account: O(log n) to find the position plus O(n) for the list
insert in the worst case. In practice that is a single memmove, so it is
fast for interview-scale n.
add_account_with_referral: the same as above, done twice (insert the new
account, then move the referrer).
get_total_credit: O(1).
top_k(k): O(k).
count_at_least(threshold): O(log n).
Space: O(n).

If list inserts are a concern, a balanced tree or sortedcontainers.SortedList
makes every write O(log n) and leaves the reads unchanged. Without ranking
queries, drop _ranked and every operation is O(1).

Test cases

add_account(50) returns 0 and add_account(80) returns 1. IDs are
sequential from 0.
Then add_account_with_referral(40, 0) returns 2, and account 0's total is
now 90.
Then add_account_with_referral(10, 2) returns 3. Account 2's total is 50,
and account 0 stays at 90 because only the direct referrer gets credit.
top_k(2) returns [0, 1] (totals 90 and 80). top_k(10) returns all
four: [0, 1, 2, 3].
count_at_least(50) returns 3 (accounts 0, 1 and 2). count_at_least(100)
returns 0.
An empty tracker: top_k(3) returns [] and count_at_least(0) returns 0.
Ties: two accounts with credit 0. top_k(2) returns [0, 1], so the lower
ID comes first.
top_k(0) returns [].
A referral to an ID that does not exist (for example 7 in a tracker with
two accounts) raises ValueError.
One account referring many: create A, then 1,000 accounts referred by A.
A's total is its own credit plus all 1,000 credits, and every write moves
only A and the new account.

Follow-ups

Multi-level referrals. A new account's credit goes to every ancestor,
perhaps with a decay per level (100%, 50%, 25%, ...). Walk the
_referrer chain and bump each ancestor. That costs O(depth) per write,
and a long chain makes it O(n). Ask whether there is a depth cap. If reads
are rare, store only own credit and compute subtree sums on demand. IDs
point to smaller IDs, so one pass from the highest ID down to 0 adds each
account's sum into its parent's.
Credit changes after creation (add_credit(id, amount) or spending).
Decide with the interviewer whether a referrer's credit should follow the
change. If it should, the change is a delta applied to the account and its
referrer, using the same _bump. If a referrer's bonus is fixed at signup,
keep a separate referral_bonus field and only change the own part.
Top k among referrers only, or ranking by the number of referrals.
Keep a second sorted index keyed by (-referrals, id) and update it in
_create. Each extra ranking is its own index, with the same update cost.
Timestamps and windows. For "top k by credit earned in the last 30
days", keep each referral event with its time. Expire events from a
queue as time moves forward, and apply negative deltas through _bump.
Scale. With millions of accounts and a stream of signups, shard the
per-account totals by ID and keep a per-shard top k. Merge the shard top-k
lists for a global answer. Totals can be eventually consistent. Make
credit additions idempotent with an event ID, so a retried signup is not
counted twice.

Interviewer expectations

Mid-level (L4)

Assigns IDs from a counter and stores state in arrays or maps indexed by ID.
Updates only the direct referrer, and checks that the referrer exists.
Gets top_k right with a sort or a heap, with a stated tie-break.

Senior (L5)

Asks early which queries will follow, and picks a structure that keeps
reads cheap without making writes expensive.
Keeps a sorted index (or explains a heap with lazy deletion) and shows
that each write moves at most two entries.
Handles edge cases without being prompted: empty tracker, k larger than
n, ties, an unknown referrer, negative credit.
Explains how multi-level credit changes the cost of a write.

Staff (L6+)

Separates the event log (signups, referrals) from derived totals, so the
totals can be rebuilt and adjustments can be replayed.
Discusses the trade-off between computing on write and on read for deep
referral chains, and uses the ID ordering to compute subtree sums in one
pass.
Covers the distributed version: sharding, merging per-shard top k,
idempotent events, and how much staleness the ranking can tolerate.