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.
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
3
4
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()
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))