改进查找最小生成树的实现

Question

I am trying to implement Kruskal's algorithm to find a minimum spanning tree in Python to solve a question on an online judge, but I am running into time limit problems. The question gives a series of edges in increasing order and asks if a minimum spanning tree is possible. Full problem specifications can be seen here.

Here is my code for the problem:

import sys
raw_input = sys.stdin.readline
line = map(int, raw_input().split())
n = line[0]
m = line[1]
dict1 = {}
lists = []

for i in xrange(1, n + 1):
    dict1[i] = set([i])

for i in xrange(m):
    edge = map(int, raw_input().split())
    a = edge[0]
    b = edge[1]
    if dict1[a] != dict1[b]:
        newSet = dict1[a].union(dict1[b])
        for vertice in newSet:
            dict1[num] = newSet
        lists.append(i + 1)

check = all(dict1[x] == dict1[1] for x in dict1.keys())

if check:
    for i in lists:
        print i
else:
    print "Disconnected Graph"

The code first creates disjoint sets for all possible vertices. Then for each edge, it checks if the sets where each of the two vertices lie are different. If they are, then the two sets are combined with a union operation. Each vertex in the combined set is a member of the newly created combined set. If the vertices are already connected, then they are skipped. I think the problem with my code is the number of times the sets have to be updated in the lines:

for vertice in newSet:
    dict1[num] = newSet

Is there a faster way to update the sets to check if they are equal? This operation is taking approximately O(vertices^2) time, and it takes too long when there are up to 100,000 vertices.

Answer 1

The key is to use the appropriate data structure for your sets. Merging sets of nodes and testing to see if any two nodes are in the same set is a classic computer science problem called "union-find".

The best data structure for this is easy and is described here:

http://www.algorithmist.com/index.php/Union_Find

Using this structure, you can merge sets and test for equality in pretty much constant time, which makes your implementation of Kruskal's algorithm (where you're provided with a pre-sorted list of edges) pretty much linear.