CSES - Connections

A computer network has nn computers numbered 1,2,,n1,2,\dots,n. The computers can have one-directional links between them. The link aba \rightarrow b means that a message can be sent from the computer aa to the computer bb.

Your task is to find if a message can be sent from any computer to any other computer after the addition of given links. A message can pass through one or more links.

In a file connections.py, implement the class Connections with the following methods:

  • add_link adds a one-directional link from the computer aa to the computer bb
  • check_network returns True if a message can be sent from any computer to any other computer, or False otherwise
class Connections:
    def __init__(self, n):
        # TODO

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

    def check_network(self):
        # TODO

if __name__ == "__main__":
    connections = Connections(5)

    connections.add_link(1, 2)
    connections.add_link(2, 3)
    connections.add_link(1, 3)
    connections.add_link(4, 5)

    print(connections.check_network()) # False

    connections.add_link(3, 5)
    connections.add_link(1, 4)

    print(connections.check_network()) # False

    connections.add_link(5, 1)

    print(connections.check_network()) # True