F - Range Connect MST Editorial by evima

Another Solution

Here is a solution using the Union-Find algorithm.

Assume \(C_1 \leq C_2 \leq \dots \leq C_Q\) (by sorting the input). Based on Kruskal’s algorithm, the answer can be obtained using the following steps:

  1. Prepare a list of sets \(\{1\}, \{2\}, \dots, \{N\}\) and initialize \(\text{ans} := 0\).
  2. For \(i = 1, 2, \dots, Q\) in this order, do the following:
    • Enumerate all sets that have at least one element between \(L_i\) and \(R_i\). Add \(C_i\) times the number of such sets to \(\text{ans}\). Then, merge those sets into one.
  3. If all sets are merged into one, output \(\text{ans}\); otherwise, output \(-1\).

This can be efficiently done using the Union-Find algorithm while maintaining the maximum vertex number of each connected component. The time complexity, including sorting, is \(O(Q(\log Q + \alpha(N)))\).

Implementation example (Python)

from atcoder.dsu import DSU

N, Q = map(int, input().split())
CLR = []
for i in range(Q):
    L, R, C = map(int, input().split())
    CLR.append((C, L - 1, R - 1))

d = DSU(N)
r = [i for i in range(N)]
ans = 0
for C, L, R in sorted(CLR):
    p = L
    while p <= R:
        p = r[d.leader(p)]
        ans += C
        if p + 1 <= R:
            nr = r[d.leader(p + 1)]
            d.merge(p, p + 1)
            r[d.leader(p)] = nr
        p += 1
print(ans if d.size(0) == N else -1)

posted:
last update: