~/problems / Weighted graphs / Union-Find

Pooled lunch money

easy ~15 min

Office workers 0 .. n-1 start with money[i] coins each. During the week, people agree to pool their money: once two people are pooled, everyone in either of their pools shares one common pot (pooling is transitive and never undone).

You get a log of events, processed in order:

  • ("pool", a, b): merge the pots of a and b (nothing changes if they already share one).
  • ("ask", a): how many coins are in a's pot right now?

Implement pot_sizes(money: list[int], events: list[tuple]) -> list[int], returning the answers to the "ask" events in order.

money = [5, 1, 7, 2]
events = [("ask", 0), ("pool", 0, 1), ("ask", 1), ("pool", 2, 3), ("pool", 1, 3), ("ask", 2), ("pool", 0, 2), ("ask", 3)]
pot_sizes(money, events)   # [5, 6, 15, 15]

Constraints: 1 <= n <= 10^5, up to 2 * 10^5 events, 0 <= money[i] <= 10^9. Asks and pools are interleaved, so you can't just group everyone at the end, and walking a whole pot for every ask is too slow.

Show hint

use union-find and keep a total per root; when two different roots merge, add the absorbed root's total to the surviving root's.

Topic: Union-Find. Path compression + union by rank; connectivity and grouping.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc