Problem 524336 · medium · Phase 05 Advanced Algorithms & Graphs

Population of Each Province

union-find · weighted merge · connected components

A country has n cities numbered 0 .. n-1; city i has population pop[i]. The list roads contains pairs [a, b] of cities joined by a two-way road. A province is a maximal set of cities that are mutually reachable by road; a city without roads is a province of its own.

Return the list of province populations (sum of the populations of the cities in each province), sorted in descending order.

Examples

Input:  pop = [5, 1, 3, 2], roads = [[0, 1], [2, 3]]
Output: [6, 5]
Explanation: cities 0 and 1 form a province of population 6; cities 2 and 3 total 5.

Input:  pop = [4, 4, 4], roads = []
Output: [4, 4, 4]

Constraints

  • 0 <= n <= 10**4, 0 <= len(roads) <= 2 * 10**4, 0 <= pop[i] <= 10**6; roads may repeat or join a city to itself
  • Target: O((n + R) * alpha(n) + P log P) time, where P is the number of provinces.

Goals

  • Maintain an aggregate (population) per union-find root
  • Enumerate roots to read off one value per component
Starting Python…