CSES - Same weight

Your task is to determine whether every spanning tree of a given graph has the same weight.

You may assume that the number of vertices is at most 50 and that the methods of the class are called at most 100 times. The weight of each edge is an integer in the range 1 \dots 1000.

In a file sameweight.py, implement a class SameWeight with the following methods:

  • constructor with the number of vertices as a parameter
  • add_edge adds an edge of weight x between vertices a and b
  • check reports whether all spanning trees have the same weight (if the graph is not connected the result should be True)
class SameWeight:
    def __init__(self,n):
        # TODO

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

    def check(self):
        # TODO

if __name__ == "__main__":
    s = SameWeight(4)
    s.add_edge(1,2,2)
    s.add_edge(1,3,3)
    print(s.check()) # True
    s.add_edge(1,4,3)
    print(s.check()) # True
    s.add_edge(3,4,3)
    print(s.check()) # True
    s.add_edge(2,4,1)
    print(s.check()) # False