Aggregate Loans by Ultimate Parent Company
Company: Affirm
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
# Aggregate Loans by Ultimate Parent Company
Each company has at most one direct parent. For every company loan, follow the parent chain to its ultimate parent and add the loan amount to that root company's total. Companies that are roots and have no children must still appear when they or a descendant have a loan. Return root totals ordered by ascending root identifier.
## Function Contract
Implement `aggregate_loans(parents, loan_companies, loan_amounts) -> list[list[int]]`. `parents[i]` is the parent of company i, or -1 when i is a root. The two loan arrays describe corresponding loans. Each output pair is `[root_id, total_amount]`.
## Constraints
- 1 <= number of companies <= 200000.
- The parent relation is a forest with no cycles.
- 0 <= number of loans <= 200000.
- 0 <= each loan amount <= 10^12.
- Every root total is guaranteed to be at most 2^53 - 1.
## Examples
```text
parents = [-1, 0, 1, -1], loan_companies = [2, 3, 1], loan_amounts = [5, 7, 4]
output = [[0, 9], [3, 7]]
```
```text
parents = [-1], loan_companies = [], loan_amounts = []
output = []
```
```hint Include a root-only borrower
Test a root company with no children that has its own loan, as well as a root with no loan anywhere in its tree.
```
```hint Exercise shared ancestry
Include loans at several depths whose companies ultimately belong to the same root, plus a separate root tree.
```
Overview: Each company has at most one direct parent. Work through the function contract, boundary cases, correctness argument, and time and space complexity expected in a production-quality solution.
Each company has at most one direct parent. For every recorded loan, follow the borrower's parent chain to its ultimate root company and add the loan amount to that root's total. Return one `[root_id, total_amount]` pair for every root whose tree contains at least one loan, sorted by ascending root identifier. A recorded zero-amount loan still causes its root to appear.
Constraints
- 1 <= len(parents) <= 200000; parents describes an acyclic forest and each value is -1 or a valid company identifier.
- 0 <= len(loan_companies) = len(loan_amounts) <= 200000; each borrower identifier is valid.
- 0 <= each loan amount <= 10^12, and every root total is at most 2^53 - 1.
Examples
Input: ([-1], [], [])
Expected Output: []
Explanation: A root with no loan anywhere in its tree contributes no row.
Input: ([-1, 0, 1, -1], [2, 3, 1], [5, 7, 4])
Expected Output: [[0, 9], [3, 7]]
Explanation: Loans at two depths share root 0 while root 3 retains its own total.
Hints
- Test a root-only borrower, a deep descendant, and borrowers in two separate root trees.
- Include repeated loans to one company and a recorded zero-amount loan.
- Exercise a large root total assembled from several individually valid loan amounts.
Community answers
Answer by akshayjagz
Python solution:
from collections import defaultdict
def find_root(company, parents, root_of):
"""Walk up to the topmost parent, caching the root for every node on the way."""
path = []
current = company
# Climb until we hit a root or a company whose root we already know.
while parents[current] != -1 and current not in root_of:
path.append(current)
current = parents[current]
root = root_of.get(current, current)
# Everyone on the path shares this root, so record it once for all of them.
for company_on_path in path:
root_of[company_on_path] = root
root_of[current] = root
return root
def aggregate_loans(parents, loan_companies, loan_amounts):
root_of = {} # company -> its topmost parent
total_by_root = defaultdict(int) # root -> summed loan amount
for company, amount in zip(loan_companies, loan_amounts):
total_by_root[find_root(company, parents, root_of)] += amount
return sorted([root, total] for root, total in total_by_root.items())