CSES - Best route

Bitland has n cities numbered 1,2,\dots,n. You know the roads between the cities, and your task is to find the shortest route between two cities.

In a file bestroute.py, implement the class BestRoute with the following methods:

  • add_road adds a road of length x between the cities a and b
  • find_route returns the length of the shortest route between two cities a and b

All the roads go both directions and thus the order of the two cities in the method add_road makes no difference. You may assume that there is at most one road between any two cities.

If there is no route, the method find_route should return -1.

class BestRoute:
    def __init__(self, n):
        # TODO

    def add_road(self, a, b, x):
        # TODO

    def find_route(self, a, b):
        # TODO

if __name__ == "__main__":
    routes = BestRoute(3)

    routes.add_road(1, 2, 2)
    print(routes.find_route(1, 3)) # -1

    routes.add_road(3, 1, 5)
    print(routes.find_route(1, 3)) # 5

    routes.add_road(2, 3, 1)
    print(routes.find_route(1, 3)) # 3