Complete the MST (Codeforces 1509F)
Tags: Greedy, dsu

As a teacher, Riko Hakozaki often needs to help her students with problems from various subjects. Today, she is asked a programming task which goes as follows.You are given an undirected complete graph with $$$n$$$ nodes, where some edges are pre-assigned with a positive weight while the rest aren't. You need to assign all unassigned edges with non-negative weights so that in the resulting fully-assigned complete graph the XOR sum of all weights would be equal to $$$0$$$.Define the ugliness of a fully-assigned complete graph the weight of its minimum spanning tree, where the weight of a spanning tree equals the sum of weights of its edges. You need to assign the weights so that the ugliness of the resulting graph is as small as possible.As a reminder, an undirected complete graph with $$$n$$$ nodes contains all edges $$$(u, v)$$$ with $$$1 \le u < v \le n$$$; such a graph has $$$\frac{n(n-1)}{2}$$$ edges.She is not sure how to solve this problem, so she asks you to solve it for her.

Input: The first line contains two integers $$$n$$$ and $$$m$$$ ($$$2 \le n \le 2 \cdot 10^5$$$, $$$0 \le m \le \min(2 \cdot 10^5, \frac{n(n-1)}{2} - 1)$$$)  — the number of nodes and the number of pre-assigned edges. The inputs are given so that there is at least one unassigned edge.The $$$i$$$-th of the following $$$m$$$ lines contains three integers $$$u_i$$$, $$$v_i$$$, and $$$w_i$$$ ($$$1 \le u_i, v_i \le n$$$, $$$u \ne v$$$, $$$1 \le w_i < 2^{30}$$$), representing the edge from $$$u_i$$$ to $$$v_i$$$ has been pre-assigned with the weight $$$w_i$$$. No edge appears in the input more than once.

Output: Print on one line one integer  — the minimum ugliness among all weight assignments with XOR sum equal to $$$0$$$.

Note: The following image showcases the first test case. The black weights are pre-assigned from the statement, the red weights are assigned by us, and the minimum spanning tree is denoted by the blue edges.  

Sample Inputs:
4 4
2 1 14
1 4 14
3 2 15
4 3 8

6 6
3 6 4
2 4 1
4 5 7
3 4 10
3 5 1
5 2 15

5 6
2 3 11
5 3 7
1 4 10
2 4 14
4 3 8
2 5 6

Sample Outputs:
15

0

6

Solution (Python):
import sys
from sys import stdin
from collections import deque

def uf_find(n, p):
    ufl = []
    while p[n] != n:
        ufl.append(n)
        n = p[n]
    for i in ufl:
        p[i] = n
    return n

def uf_union(a, b, p, rank):
    ap = uf_find(a, p)
    bp = uf_find(b, p)
    if ap == bp:
        return True
    else:
        if rank[ap] > rank[bp]:
            p[bp] = ap
        elif rank[ap] < rank[bp]:
            p[ap] = bp
        else:
            p[bp] = ap
            rank[ap] += 1
        return False
(n, m) = map(int, input().split())
if n <= 1000:
    lis = [[0] * n for i in range(n)]
    wuv = []
    XOR = 0
    for i in range(m):
        (u, v, w) = map(int, input().split())
        u -= 1
        v -= 1
        wuv.append((w, u, v))
        lis[u][v] = lis[v][u] = w
        XOR ^= w
    wuv.sort()
    p = [i for i in range(n)]
    rank = [0] * n
    edge = []
    flag = False
    for i in range(n):
        for j in range(i):
            if lis[i][j] == 0:
                if uf_union(i, j, p, rank):
                    flag = True
                else:
                    edge.append((i, j, 0))
    nans = 0
    for (w, u, v) in wuv:
        if uf_union(u, v, p, rank):
            pass
        else:
            nans += w
            edge.append((u, v, w))
    if flag:
        print(nans)
        sys.exit()
    p = [i for i in range(n)]
    rank = [0] * n
    for (i, j, w) in edge:
        if w > 0:
            uf_union(i, j, p, rank)
    nplus = XOR
    for (w, u, v) in wuv:
        if uf_find(u, p) != uf_find(v, p):
            nplus = min(nplus, w)
            break
        elif w > nplus:
            break
    print(nans + nplus)
else:
    lis = [[] for i in range(n)]
    wuv = []
    for i in range(m):
        (u, v, w) = map(int, stdin.readline().split())
        u -= 1
        v -= 1
        wuv.append((w, u, v))
        lis[u].append((v, w))
        lis[v].append((u, w))
    able = [True] * n
    p = [i for i in range(n)]
    rank = [0] * n
    for st in range(n):
        if not able[st]:
            continue
        block = [0] * n
        cnt = 0
        able[st] = False
        q = deque([st])
        while q:
            while q:
                nv = q.popleft()
                cnt += 1
                for (u, w) in lis[nv]:
                    if able[u] and block[u] == cnt - 1:
                        block[u] = cnt
            for i in range(n):
                if able[i] and block[i] != cnt:
                    uf_union(i, st, p, rank)
                    able[i] = False
                    q.append(i)
    wuv.sort()
    print(sum([w for (w, u, v) in wuv if not uf_union(u, v, p, rank)]))

Solution (Python):
import sys
from sys import stdin
from collections import deque

def uf_find(n, p):
    ufl = []
    while p[n] != n:
        ufl.append(n)
        n = p[n]
    for i in ufl:
        p[i] = n
    return n

def uf_union(a, b, p, rank):
    ap = uf_find(a, p)
    bp = uf_find(b, p)
    if ap == bp:
        return True
    else:
        if rank[ap] > rank[bp]:
            p[bp] = ap
        elif rank[ap] < rank[bp]:
            p[ap] = bp
        else:
            p[bp] = ap
            rank[ap] += 1
        return False
(n, m) = map(int, stdin.readline().split())
if n <= 1000:
    lis = [[0] * n for i in range(n)]
    wuv = []
    XOR = 0
    for i in range(m):
        (u, v, w) = map(int, stdin.readline().split())
        u -= 1
        v -= 1
        wuv.append((w, u, v))
        lis[u][v] = lis[v][u] = w
        XOR ^= w
    wuv.sort()
    p = [i for i in range(n)]
    rank = [0] * n
    edge = []
    flag = False
    for i in range(n):
        for j in range(i):
            if lis[i][j] == 0:
                if uf_union(i, j, p, rank):
                    flag = True
                else:
                    edge.append((i, j, 0))
    nans = 0
    for (w, u, v) in wuv:
        if uf_union(u, v, p, rank):
            pass
        else:
            nans += w
            edge.append((u, v, w))
    if flag:
        print(nans)
        sys.exit()
    p = [i for i in range(n)]
    rank = [0] * n
    for (i, j, w) in edge:
        if w > 0:
            uf_union(i, j, p, rank)
    nplus = XOR
    for (w, u, v) in wuv:
        if uf_find(u, p) != uf_find(v, p):
            nplus = min(nplus, w)
            break
        elif w > nplus:
            break
    print(nans + nplus)
else:
    lis = [[] for i in range(n)]
    wuv = []
    for i in range(m):
        (u, v, w) = map(int, stdin.readline().split())
        u -= 1
        v -= 1
        wuv.append((w, u, v))
        lis[u].append((v, w))
        lis[v].append((u, w))
    able = [True] * n
    p = [i for i in range(n)]
    rank = [0] * n
    for st in range(n):
        if not able[st]:
            continue
        block = [0] * n
        cnt = 0
        able[st] = False
        q = deque([st])
        while q:
            while q:
                nv = q.popleft()
                cnt += 1
                for (u, w) in lis[nv]:
                    if able[u] and block[u] == cnt - 1:
                        block[u] = cnt
            for i in range(n):
                if able[i] and block[i] != cnt:
                    uf_union(i, st, p, rank)
                    able[i] = False
                    q.append(i)
    wuv.sort()
    ans = 0
    for (w, u, v) in wuv:
        if not uf_union(u, v, p, rank):
            ans += w
    print(ans)

Solution (Python):
from collections import deque

def uf_find(n, p):
    ufl = []
    while p[n] != n:
        ufl.append(n)
        n = p[n]
    for i in ufl:
        p[i] = n
    return n

def uf_union(a, b, p, rank):
    ap = uf_find(a, p)
    bp = uf_find(b, p)
    if ap == bp:
        return True
    else:
        if rank[ap] > rank[bp]:
            p[bp] = ap
        elif rank[ap] < rank[bp]:
            p[ap] = bp
        else:
            p[bp] = ap
            rank[ap] += 1
        return False
(n, m) = map(int, input().split())
if n <= 1000:
    lis = [[0] * n for i in range(n)]
    wuv = []
    XOR = 0
    for i in range(m):
        (u, v, w) = map(int, input().split())
        u -= 1
        v -= 1
        wuv.append((w, u, v))
        lis[u][v] = lis[v][u] = w
        XOR ^= w
    wuv.sort()
    p = [i for i in range(n)]
    rank = [0] * n
    edge = []
    flag = False
    nans = 0
    for i in range(n):
        for j in range(i):
            if lis[i][j] == 0:
                if uf_union(i, j, p, rank):
                    flag = True
                else:
                    edge.append((i, j, 0))
    for (w, u, v) in wuv:
        if uf_union(u, v, p, rank):
            pass
        else:
            nans += w
            edge.append((u, v, w))
    if flag:
        print(nans)
        exit()
    p = [i for i in range(n)]
    rank = [0] * n
    for (i, j, w) in edge:
        if w > 0:
            uf_union(i, j, p, rank)
    nplus = XOR
    for (w, u, v) in wuv:
        if uf_find(u, p) != uf_find(v, p):
            nplus = min(nplus, w)
            break
        elif w > nplus:
            break
    print(nans + nplus)
else:
    lis = [[] for i in range(n)]
    wuv = []
    for i in range(m):
        (u, v, w) = map(int, input().split())
        u -= 1
        v -= 1
        wuv.append((w, u, v))
        lis[u].append((v, w))
        lis[v].append((u, w))
    able = [True] * n
    p = [i for i in range(n)]
    rank = [0] * n
    for st in range(n):
        if not able[st]:
            continue
        block = [0] * n
        cnt = 0
        able[st] = False
        q = deque([st])
        while q:
            while q:
                nv = q.popleft()
                cnt += 1
                for (u, w) in lis[nv]:
                    if able[u] and block[u] == cnt - 1:
                        block[u] = cnt
            for i in range(n):
                if able[i] and block[i] != cnt:
                    uf_union(i, st, p, rank)
                    able[i] = False
                    q.append(i)
    wuv.sort()
    print(sum([w for (w, u, v) in wuv if not uf_union(u, v, p, rank)]))