https://www.hackerrank.com/challenges/journey-to-the-moon

moon.py

from collections import defaultdict

import itertools


def data():
    N, I = [int(i) for i in input().split()]
    rows = [[int(i) for i in input().split()] for _ in range(I)]
    return N, rows


def ev(N, rows):
    gr_i = 0
    groups = defaultdict(set)
    m = [None] * N

    for A, B in rows:
        if m[A] and m[B]:
            if m[A] != m[B]:
                mi, ma = min(m[A], m[B]), max(m[A], m[B])
                for k in groups[mi]:
                    groups[ma].add(k)
                    m[k] = ma
                del groups[mi]
        else:
            if m[A]:
                gr = m[A]
            elif m[B]:
                gr = m[B]
            else:
                gr_i += 1
                gr = gr_i
            groups[gr].add(A)
            m[A], m[B] = gr, gr
            groups[gr].add(B)
    for i in range(N):
        if m[i] is None:
            gr_i += 1
            gr = gr_i
            groups[gr].add(i)
            m[i] = gr
    s = [len(i) for i in groups.values()]
    print(sum(a*b for a, b in itertools.combinations(s, 2)))


ev(*data())