MST challenge gives "time exceeded" error

Question

I am doing the BLINNET problem on Sphere Online Judge where I need to find the cost of a minimum spanning tree. I should follow a structure with Edge and Vertex instances. Vertices represent cities in this case.

I get a "time exceeded" error, and I feel like too many for loop iterations are at the cause, but that is the best I can do. I want to try the binary sort to see if it works with that, but that is not easy as it should be sorted using the key property in the City class.

Sample input

2

4
gdansk
2
2 1
3 3
bydgoszcz
3
1 1
3 1
4 4
torun
3
1 3
2 1
4 1
warszawa
2
2 4
3 1

3
ixowo
2
2 1
3 3
iyekowo
2
1 1
3 7
zetowo
2
1 3
2 7

Output for the Sample

3
4

My code

import sys
import heapq

class City:

    def __init__(self, city_id):
        self.city_id = city_id
        self.key = float('inf')
        self.parent = None
        self.edge_list = list()
        self.visited = False
        #self.city_name = None

    def is_not_visited(self):
        if self.visited is False:
            return True
        return False

    def add_neighbor(self, edge):
        self.edge_list.append(edge)

    def __lt__(self, other):
        return self.key < other.key

class Edge:

    def __init__(self, to_vertex, cost):
        self.to_vertex = to_vertex
        self.cost = cost

#
# def find_and_pop(queue):
#     min = queue[0]
#     index = 0
#     for a in range(0, len(queue)):
#         if queue[a].key < min.key:
#             min = queue[a]
#             index = a
#     return queue.pop(index)
#

def MST(vertices_list):
    queue = vertices_list
    current = queue[0]
    current.key = 0
    #visited_list = list()
    #heapq.heapify(queue)
    total_weight = 0
    while queue:
        #current = find_and_pop(queue)
        current = queue.pop(0)
        for edge in current.edge_list:
            if edge.to_vertex.is_not_visited():
                if edge.cost < edge.to_vertex.key:
                    edge.to_vertex.key = edge.cost
                    edge.to_vertex.parent = current
        total_weight = total_weight + current.key
        current.visited = True
        queue = sorted(queue, key=lambda x: x.city_id)
        #heapq.heapify(queue)
        #visited_list.append(current)

    # total_weight = 0
    # for x in visited_list:
    #     total_weight = total_weight + x.key
    sys.stdout.write("{0}\n".format(total_weight))

class TestCase:
    def __init__(self, vertices):
        self.vertices = vertices

testcases = []

def main():

    case_num = int(sys.stdin.readline())
    #skip_line = sys.stdin.readline()
    for n_case in range(0, case_num):
        sys.stdin.readline()
        vertices_list = list()
        number_of_city = int(sys.stdin.readline())
        #interate and make for the time of number of cities
        for n_city in range(0, number_of_city):
            city = City(n_city)
            vertices_list.append(city)

        for n_city in range(0, number_of_city):
            c_name = sys.stdin.readline()
            #vertices_list[n_city].city_name = c_name
            num_neighbor = int(sys.stdin.readline())
            for n_neigh in range(0, num_neighbor):
                to_city_cost = sys.stdin.readline()
                to_city_cost = to_city_cost.split(" ")
                to_city = int(to_city_cost[0])
                cost = int(to_city_cost[1])
                edge = Edge(vertices_list[to_city-1], cost)
                vertices_list[n_city].edge_list.append(edge)

        testcase = TestCase(vertices_list)
        testcases.append(testcase)

    count = 0
    for testcase in testcases:
        MST(testcase.vertices)
        # if count < case_num -1:
        #     print()
        # count = count + 1



if __name__ == "__main__":
    main()

Answer 1

The sorted call in your MST loop makes the solution inefficient. You have some commented-out code that relies on heapq, and that is indeed the way to avoid having to sort the queue each time you alter it. Anyway, I don't understand why you would sort the queue by city id. If anything, it should be sorted by key.

Although it could work with the key property as you did it, it seems more natural to me to add edges to the queue (heap) instead of vertices, so you have the edge cost as the basis for the heap property. Also, that queue should not have all the items from the start, but add them as they are selected during the algorithm. And, that corresponds more the the MST-building algorithm, which adds edge after edge, each time the one with the minimum cost.

If edges are pushed on a heap, they must be comparable. So __lt__ must be implemented on the Edge class like you did for the Vertex class.

class Edge:
    # ... your code remains unchanged... Just add:
    def __lt__(self, other):
        return self.cost < other.cost

def MST(vertices_list):
    # first edge in the queue is a virtual one with zero cost.
    queue = [Edge(vertices_list[0], 0)] # heap of edges, ordered by cost
    total_weight = 0
    while queue:
        mst_edge = heapq.heappop(queue) # pop both cost & vertex
        current = mst_edge.to_vertex
        if current.visited: continue
        for edge in current.edge_list:
            if not edge.to_vertex.visited:
                heapq.heappush(queue, edge)
        current.visited = True
        total_weight += mst_edge.cost

    sys.stdout.write("{0}\n".format(total_weight))