Coding sample solution

Coding: optimization over choices

How to talk through this coding question, from a first idea to working code, in an interview.

Coding interview · Sample solution in Python · Independent, not affiliated with Uber

The coding question

You're given a company's org chart and a fun score for each employee. You can't invite an employee and their direct manager to the same party. Return the most total fun you can get and who you'd invite.

Ask a few questions first

AskWhat we'll assume
How's the org chart given?As a list managers, where managers[i] is employee i's direct manager, and -1 for the one person at the top.
Can fun scores be negative?No, they're zero or more.
Is it only the direct manager that's off limits?Yes. You can invite someone and their manager's manager.
If two guest lists tie, which one?Any of them. Return the employee numbers sorted.

A couple of examples

InputOutput
managers [-1, 0, 0, 1, 1, 2], fun [3, 4, 5, 6, 7, 1](18, [2, 3, 4])
managers [-1, 0, 0], fun [10, 3, 4](10, [0])
managers [-1, 0, 1, 2], fun [1, 0, 0, 1](2, [0, 3])

Start with brute force

Try every subset of employees, throw out the ones that put someone with their direct manager, and keep the subset with the most fun.

Cost: 2^n subsets, each checked in O(n). Fine for 15 people, hopeless for 1,000.

A better way

Look at one manager and their team underneath. Either the manager comes, and then none of their direct reports can, or they don't, and then each report is free to come or not. Either way, the best for each report's team doesn't depend on anything above the manager. That's dynamic programming on a tree.

So work out two numbers for everyone: invite, the best fun in their team if they come, and skip, the best if they don't. invite is their own fun plus each report's skip. skip is the sum of each report's better option. Handle reports before their manager, and the boss's better number is the answer.

  1. Build a list of direct reports for each employee and find the boss.
  2. Make a top-down order: start with the boss and keep appending each person's reports.
  3. Go through that order backwards, so reports are done before their manager, and fill in invite and skip.
  4. To get the guest list, go top-down. If someone's manager is coming, they can't. Otherwise, invite them if invite beats skip.
  5. Return the boss's better number and the sorted guest list.

The code

solution.pypython
def best_party(managers, fun):
    n = len(fun)
    reports = [[] for _ in range(n)]
    boss = -1
    for person, manager in enumerate(managers):
        if manager == -1:
            boss = person
        else:
            reports[manager].append(person)

    # Visit everyone top-down, then work through that list backwards,
    # so each person is handled after all their reports.
    order = [boss]
    for person in order:
        order.extend(reports[person])

    invite = [0] * n  # best fun in this person's team if they come
    skip = [0] * n    # best fun in this person's team if they don't
    for person in reversed(order):
        invite[person] = fun[person]
        for r in reports[person]:
            invite[person] += skip[r]
            skip[person] += max(invite[r], skip[r])

    # Walk back down to see who we picked.
    guests = []
    manager_coming = [False] * n
    for person in order:
        boss_of = managers[person]
        if boss_of != -1 and manager_coming[boss_of]:
            continue
        if invite[person] > skip[person]:
            manager_coming[person] = True
            guests.append(person)
    return max(invite[boss], skip[boss]), sorted(guests)
CostWhy
TimeO(n log n)Each employee and each manager link is handled a few times. Sorting the guest list is the log n
SpaceO(n)The report lists, the order and the two numbers per person

Test it

Walk through one of the examples out loud, then check the cases that tend to break things:

  • Just the boss: invite is their fun, so they come.
  • A long chain, like everyone managing one person: the order list handles it without recursion, so there's no depth limit to hit.
  • A manager whose reports are worth more together, like 10 against 4 and 5 and 7: skip wins and the manager stays home.
  • Ties between invite and skip: we skip the person, which is still a best answer.
  • All zeros: the total is 0 and nobody needs to come.

Common mistakes

  • Greedily inviting the most fun people first. One star can block a whole team that adds up to more.
  • Thinking it's every other level, like inviting all managers or all their reports. A team can skip two levels in a row.
  • Writing it recursively and hitting Python's recursion limit on a deep chart.
  • Picking guests bottom-up without checking whether their manager ended up coming.

Follow-up questions

What if the boss has to come?

Return invite of the boss instead of the better of the two, and start the guest walk by inviting them.

What if you can't invite someone and their manager's manager either?

Keep more history. Invite now adds the best from two levels down, so you need a third number per person, like 'not coming, and no reports coming either'.

What if it's a list instead of a tree, like a driver who can't work two days in a row?

Same idea with one report each: take is the best if they work today, skip is the best if they don't, and you only need yesterday's two numbers.

How the time adds up

MinutesWhat you're doing
0–5Ask questions and agree on examples
5–10Brute force and what it costs
10–15Explain the better idea
15–30Write the code
30–38Test it out loud
38–45Follow-up questions

Key concepts used here

We wrote this sample solution ourselves. Practice Loop is independent and isn't affiliated with or endorsed by Uber.