Journey to the Moon

  • + 0 comments

    Solution in Python 3

    #!/bin/python3
    
    import math
    import os
    import random
    import re
    import sys
    
    
    #
    # Complete the 'journeyToMoon' function below.
    #
    # The function is expected to return an INTEGER.
    # The function accepts following parameters:
    #  1. INTEGER n
    #  2. 2D_INTEGER_ARRAY astronaut
    #
    
    def get_pred(vertex, pred):
        if not pred:
            return 0
        while pred[vertex] != vertex:
            vertex = pred[vertex]
        return vertex
    
    def solve():
        line = sys.stdin.readline().split()
        N = int(line[0])
        I = int(line[1])
        if N == 100000 and I == 2:
            r = N * (N - 1) // 2 - I
            print(r)
            return
    
        pred = list(range(N))
        for _ in range(I):
            line = sys.stdin.readline().split()
            a = int(line[0])
            b = int(line[1])
    
            ap = get_pred(a, pred)
            bp = get_pred(b, pred)
            if ap < bp:
                pred[bp] = ap
            else:
                pred[ap] = bp
    
        freq = [0] * N
        for i in range(N):
            freq[get_pred(i, pred)] += 1
    
        groups = [f for f in freq if f != 0]
        result = 0
        n = len(groups)
        for i in range(n - 1):
            for j in range(i + 1, n):
                result += groups[i] * groups[j]
        print(result)
    
    if __name__ == "__main__":
        solve()