CODEWITHSUNDEEP
pythonalgorithmgraph-theoryminimum-spanning-treekruskals-algorithm
Bob Marshall
Bob Marshall

Reputation: 453

Improving Implementation of Finding a Minimum Spanning Tree

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.

Upvotes: 1

Views: 1065

Answers (1)

Matt Timmermans
Matt Timmermans

Reputation: 59144

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.

Upvotes: 3

Related Questions