CSES - Same weigth

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

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

  • add_edge adds an edge of weight xx between the nodes aa and bb
  • check reports whether all spanning trees have the same weight

If the graph is not connected, the method check should return True, because then there are no spanning trees.

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

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

    def check(self):
        # TODO

if __name__ == "__main__":
    same_weight = SameWeight(4)

    same_weight.add_edge(1, 2, 2)
    same_weight.add_edge(1, 3, 3)
    print(same_weight.check()) # True

    same_weight.add_edge(1, 4, 3)
    print(same_weight.check()) # True

    same_weight.add_edge(3, 4, 3)
    print(same_weight.check()) # True

    same_weight.add_edge(2, 4, 1)
    print(same_weight.check()) # False