[hackerrank][python]Journey to the Moon

hackerrank: Journey to the Moon

1st attempt

# Determine how many pairs of astronauts from different countries they can choose from.
# UN 연합은 달에 사람들을 보내려고 한다.
# 그들은 이 사람들이 서로 다른 나라였으면 한다.
# 당신은 두 쌍의 우주조종사 id 리스트를 받을 것이다.
# 각 쌍은 같은 나라 우주조종사 id의 두 쌍으로 이루어져 있다.
# - 얼마나 많은 조종사 쌍을 만들 수 있는지 말해라.

# n : number of astronaut
# p : the number of pairs
from collections import defaultdict

# 일단 연결된 것들을 만든다.
def dfs(graph, start):
  visited, stack = set(), [start]
  while stack:
    v = stack.pop()
    visited.add(v)
    # 해당 정점에 연결된 것들을 stack에 넣는다.
    # 방문한 것은 제외
    stack.extend(graph[v] - visited)
  # 연결된 것들을 리턴.
  return visited

def journeyToMoon(n, ast_pairs):
  graph = defaultdict(set)
  asts = set()
  # 그래프 연결...
  for a1, a2 in ast_pairs:
    graph[a1].add(a2)
    graph[a2].add(a1)
    asts.add(a1)
    asts.add(a2)
  
  # 고립된 애들이 있을 수 있나? 일단 받아놔... (있는듯...)
  iso_asts = set([i for i in range(n)]) - asts
  
  groups = []  # divied by contry
  while asts:
    ast = asts.pop()
    connected_asts = dfs(graph, ast)
    groups.append(connected_asts)
    # 한번 방문한 팀은 다시 방문하지 않는다.
    asts = asts - connected_asts
  
  for i in iso_asts:
    groups.append({i})
  
  idx1, idx2 = 0, 1
  group_sum = 0
  while idx1 < len(groups)-1:
    idx2 = idx1 + 1
    while idx2 < len(groups):
      group_sum += len(groups[idx1]) * len(groups[idx2])
      idx2 += 1
    idx1 += 1
 
  return group_sum
    
  
n, p = list(map(int, input().strip().split(' ')))
ast_pairs = []
for _ in range(p):
   astronaut = [int(a) for a in input().strip().split(' ')]
   ast_pairs.append(astronaut)
   
result = journeyToMoon(n, ast_pairs)
print(result)

def dfs(graph, start):
  visited, stack = set(), [start]
  while stack:
    v = stack.pop()
    visited.add(v)
    stack.extend(graph[v] - visited)
  return visited

def journeyToMoon(n, ast_pairs):
  graph = defaultdict(set)
  asts = set()
  for a1, a2 in ast_pairs:
    graph[a1].add(a2)
    graph[a2].add(a1)
    asts.add(a1)
    asts.add(a2)
  
  iso_asts = set([i for i in range(n)]) - asts
  
  groups = []  # divied by contry
  while asts:
    ast = asts.pop()
    connected_asts = dfs(graph, ast)
    groups.append(connected_asts)
    # 한번 방문한 팀은 다시 방문하지 않는다.
    asts = asts - connected_asts
  
  for i in iso_asts:
    groups.append({i})
  
  # 뒤에 추가된 것들을 따로 분리해야 한다. 그래야 추가된 녀석들이 만드는
  # 패턴을 알 수 있다.
  # A*B + A*C + B*C -- A*B + A*C + B*C -- A*B + (A+B)*C 
  # A*B + A*C + A*D + B*C + B*D + C*D 
  # - A*B + (A+B)*C (A+B+C)*D
  # A*B + A*C + A*D + A*E + B*C + B*D + B*E + C*D + C*E + D*E
  # - A*B + (A+B)*C + (A+B+C)*D + (A+B+C+D)*E 
  idx1, idx2 = 0, 1
  group_sum = 0
  tmp_sum = 0
  while idx1 < len(groups):
    group_sum += tmp_sum*len(groups[idx1])
    tmp_sum += len(groups[idx1])

    # group_sum += tmp_sum*(len(groups[idx1 + 1]))
    idx1 += 1
  
  return group_sum
    
  
n, p = list(map(int, input().strip().split(' ')))
ast_pairs = []
for _ in range(p):
   astronaut = [int(a) for a in input().strip().split(' ')]
   ast_pairs.append(astronaut)
   
result = journeyToMoon(n, ast_pairs)
print(result)