Avoid Rainbow Cycles (Codeforces 1408E)
Tags: Greedy, dsu

You are given $$$m$$$ sets of integers $$$A_1, A_2, \ldots, A_m$$$; elements of these sets are integers between $$$1$$$ and $$$n$$$, inclusive.There are two arrays of positive integers $$$a_1, a_2, \ldots, a_m$$$ and $$$b_1, b_2, \ldots, b_n$$$. In one operation you can delete an element $$$j$$$ from the set $$$A_i$$$ and pay $$$a_i + b_j$$$ coins for that.You can make several (maybe none) operations (some sets can become empty).After that, you will make an edge-colored undirected graph consisting of $$$n$$$ vertices. For each set $$$A_i$$$ you will add an edge $$$(x, y)$$$ with color $$$i$$$ for all $$$x, y \in A_i$$$ and $$$x < y$$$. Some pairs of vertices can be connected with more than one edge, but such edges have different colors.You call a cycle $$$i_1 \to e_1 \to i_2 \to e_2 \to \ldots \to i_k \to e_k \to i_1$$$ ($$$e_j$$$ is some edge connecting vertices $$$i_j$$$ and $$$i_{j+1}$$$ in this graph) rainbow if all edges on it have different colors.Find the minimum number of coins you should pay to get a graph without rainbow cycles.

Input: The first line contains two integers $$$m$$$ and $$$n$$$ ($$$1 \leq m, n \leq 10^5$$$), the number of sets and the number of vertices in the graph.The second line contains $$$m$$$ integers $$$a_1, a_2, \ldots, a_m$$$ ($$$1 \leq a_i \leq 10^9$$$).The third line contains $$$n$$$ integers $$$b_1, b_2, \ldots, b_n$$$ ($$$1 \leq b_i \leq 10^9$$$).In the each of the next of $$$m$$$ lines there are descriptions of sets. In the $$$i$$$-th line the first integer $$$s_i$$$ ($$$1 \leq s_i \leq n$$$) is equal to the size of $$$A_i$$$. Then $$$s_i$$$ integers follow: the elements of the set $$$A_i$$$. These integers are from $$$1$$$ to $$$n$$$ and distinct.It is guaranteed that the sum of $$$s_i$$$ for all $$$1 \leq i \leq m$$$ does not exceed $$$2 \cdot 10^5$$$.

Output: Print one integer: the minimum number of coins you should pay for operations to avoid rainbow cycles in the obtained graph.

Note: In the first test, you can make such operations:  Delete element $$$1$$$ from set $$$1$$$. You should pay $$$a_1 + b_1 = 5$$$ coins for that.  Delete element $$$1$$$ from set $$$2$$$. You should pay $$$a_2 + b_1 = 6$$$ coins for that. You pay $$$11$$$ coins in total. After these operations, the first and the second sets will be equal to $$$\{2\}$$$ and the third set will be equal to $$$\{1, 2\}$$$.So, the graph will consist of one edge $$$(1, 2)$$$ of color $$$3$$$.In the second test, you can make such operations:  Delete element $$$1$$$ from set $$$1$$$. You should pay $$$a_1 + b_1 = 11$$$ coins for that.  Delete element $$$4$$$ from set $$$2$$$. You should pay $$$a_2 + b_4 = 13$$$ coins for that.  Delete element $$$7$$$ from set $$$3$$$. You should pay $$$a_3 + b_7 = 13$$$ coins for that.  Delete element $$$4$$$ from set $$$4$$$. You should pay $$$a_4 + b_4 = 16$$$ coins for that.  Delete element $$$7$$$ from set $$$6$$$. You should pay $$$a_6 + b_7 = 13$$$ coins for that. You pay $$$66$$$ coins in total.After these operations, the sets will be:  $$$\{2, 3\}$$$;  $$$\{1\}$$$;  $$$\{1, 3\}$$$;  $$$\{3\}$$$;  $$$\{3, 4, 5, 6, 7\}$$$;  $$$\{5\}$$$;  $$$\{8\}$$$. We will get the graph:There are no rainbow cycles in it.

Sample Inputs:
3 2
1 2 3
4 5
2 1 2
2 1 2
2 1 2

7 8
3 6 7 9 10 7 239
8 1 9 7 10 2 6 239
3 2 1 3
2 4 1
3 1 3 7
2 4 3
5 3 4 5 6 7
2 5 7
1 8

Sample Outputs:
11

66

Solution (Python):
import sys
input = sys.stdin.buffer.readline

def getlist():
    return list(map(int, input().split()))

class UnionFind:

    def __init__(self, N):
        self.par = [i for i in range(N)]
        self.rank = [0] * N

    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def same_check(self, x, y):
        return self.find(x) == self.find(y)

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

class Kruskal:

    def __init__(self, E, N):
        self.uf = UnionFind(N)
        self.cost = 0
        for (w, u, v) in E:
            if self.uf.same_check(u, v) != True:
                self.uf.union(u, v)
                self.cost += w

    def mincost(self):
        return self.cost

def main():
    (M, N) = getlist()
    A = getlist()
    B = getlist()
    E = []
    MAX = 0
    for i in range(M):
        num = getlist()
        s = num[0]
        for j in range(1, s + 1):
            bj = num[j] - 1
            w = A[i] + B[bj]
            E.append((w, i, M + bj))
            MAX += w
    E.sort(key=lambda x: x[0], reverse=True)
    K = Kruskal(E, N + M)
    res = K.mincost()
    ans = MAX - res
    print(ans)
if __name__ == '__main__':
    main()

Solution (Python):
import os
import sys
from io import BytesIO
input = BytesIO(os.read(0, os.fstat(0).st_size)).readline

def root(x):
    if x == p[x]:
        return x
    p[x] = y = root(p[x])
    return y

def unite(x, y):
    px = root(x)
    py = root(y)
    if px == py:
        return 0
    rx = rank[px]
    ry = rank[py]
    if ry < rx:
        p[py] = px
    elif rx < ry:
        p[px] = py
    else:
        p[py] = px
        rank[px] += 1
    return 1
if __name__ == '__main__':
    (M, N) = map(int, input().split())
    A = list(map(int, input().split()))
    B = list(map(int, input().split()))
    E = []
    R = 0
    for i in range(M):
        X = list(map(int, input().split()))
        for j in X[1:]:
            E.append((N + i, j - 1, -A[i] - B[j - 1]))
            R += A[i] + B[j - 1]
    E.sort(key=lambda x: x[2])
    (*p,) = range(N + M)
    rank = [1] * (N + M)
    for (v, u, c) in E:
        if unite(v, u):
            R += c
    print(R)

Solution (Python):
import sys
input = sys.stdin.buffer.readline

def getlist():
    return list(map(int, input().split()))

class UnionFind:

    def __init__(self, N):
        self.par = [i for i in range(N)]
        self.rank = [0] * N

    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def same_check(self, x, y):
        return self.find(x) == self.find(y)

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

class Kruskal:

    def __init__(self, E, N):
        self.uf = UnionFind(N)
        self.cost = 0
        for i in range(len(E)):
            (w, u, v) = E[i]
            if self.uf.same_check(u, v) != True:
                self.uf.union(u, v)
                self.cost += w

    def maxcost(self):
        return self.cost

def main():
    (M, N) = getlist()
    A = getlist()
    B = getlist()
    E = []
    MAX = 0
    for i in range(M):
        num = getlist()
        s = num[0]
        for j in range(1, s + 1):
            bj = num[j] - 1
            w = A[i] + B[bj]
            E.append((w, i, M + bj))
            MAX += w
    E.sort(key=lambda x: x[0], reverse=True)
    K = Kruskal(E, N + M)
    res = K.maxcost()
    ans = MAX - res
    print(ans)
if __name__ == '__main__':
    main()