[HackerRank][Python3] Journey to the Moon

Problem :

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

My Solution :

#!/usr/bin/env python3

def journeyToMoon(n, astronaut):
    country = {x: 0 for x in range(n)}
    c_code = 1
    for a1, a2 in astronaut:
        c1, c2 = country[a1], country[a2]
        if c1 == c2 == 0:
            country[a1], country[a2] = c_code, c_code
            c_code += 1
        elif c1 * c2 == 0:
            c = max(c1, c2)
            country[a1], country[a2] = c, c
        elif c1 != c2:
            for a, c in country.items():
                if c == c2:
                    country[a] = c1
            
    country_count = {}
    for c in country.values():
        country_count.setdefault(c, 0)
        country_count[c] += 1
    
    sum, res = 0, 0
    for c in country_count.values():
        res += sum * c
        sum += c
        
    singles = country_count.get(0, 0)
    if singles > 1:
        res += singles*(singles-1)//2
    
    return res


n, p = map(int, input().strip().split())
astronaut = []
for _ in range(p):
    astronaut.append(list(map(int, input().strip().split())))
result = journeyToMoon(n, astronaut)
print(result)

My Solution2:

#!/usr/bin/env python3

def find_set_index(a):
    for i in range(len(a_set_list)):
        if a in a_set_list[i]:
            return i

def journeyToMoon(n, astronaut):
    for a1, a2 in astronaut:
        i1, i2 = find_set_index(a1), find_set_index(a2)
        if i1 != i2:
            a_set_list[i1] = a_set_list[i1].union(a_set_list[i2])
            del a_set_list[i2]
    
    sum, res = 0, 0
    for s in a_set_list:
        res += sum * len(s)
        sum += len(s)
    
    return res
        

n, p = map(int, input().strip().split())
astronaut = []
for _ in range(p):
    astronaut.append(list(map(int, input().strip().split())))
a_set_list = [{x} for x in range(n)]    
result = journeyToMoon(n, astronaut)
print(result)