[해커랭크] Journey to the Moon

문제 링크

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

Problem

The member states of the UN are planning to send people to the moon. They want them to be from different countries. You will be given a list of pairs of astronaut ID's. Each pair is made of astronauts from the same country. Determine how many pairs of astronauts from different countries they can choose from.

유엔 회원국들은 달에 사람들을 보낼 계획이다. 그들은 그들이 다른 나라 출신이기를 원한다. 우주 비행사 ID의 한 쌍 목록이 주어질 것이다. 각각의 한 쌍은 같은 나라 출신의 우주 비행사들로 이루어져 있다. 다른 나라의 우주비행사 몇 쌍을 선택할 수 있는지 결정하세요.

문제는 파파고가 번역해주었다.

Example

n = 4
astronaut = [1, 2], [2, 3]
0, 1, 2, 3까지의 번호를 가진 4명의 우주비행사가 있다. [0] 이 같은 국가, [1, 2, 3]이 같은 국가이다.
[0, 1], [0, 2], [0, 3] 과 같은 조합으로 우주에 갈 수 있다.

틀린 풀이

상당히 거칠게 풀이하였고 틀렸기 때문에 나같이 푸는 사람은 없길..

import math
import os
import random
import re
import sys
from itertools import combinations

def journeyToMoon(n, astronaut):
	# 같은 국가가 없는 경우 -1
    k = dict()
    for _ in range(0, n):
        k[_] = -1

    # 같은 국가가 있는 지 확인
    # 다음 작업을 하였을 때 나오는 결과물
    # k = {0: -1, 1: 0, 2: 0, 3: 0}
    for i, a in enumerate(astronaut):
        if k[a[0]] == -1 and k[a[1]] != -1:
            k[a[0]] = k[a[1]]
        elif k[a[0]] != -1 and k[a[1]] == -1:
            k[a[1]] = k[a[0]]
        elif k[a[0]] == -1 and k[a[1]] == -1:
            k[a[0]], k[a[1]] = i, i
        # 각각이 같은 국가가 있는데, 각각의 수가 다른 경우
        # 더 작은 수로 국가를 통일하기 위해서 하는 작업
        else:
            little_min = 0
            _min = 0
            if k[a[0]] > k[a[1]]:
                little_min = k[a[0]]
                _min = k[a[1]]
            else:
                little_min = k[a[1]]
                _min = k[a[0]]
   
            for i in range(n):
                if k[i] == little_min:
                    k[i] = _min
     
    # 다음 작업을 하였을 때 나오는 결과물
    # nation:  {-1: [0], 0: [1, 2, 3]}
    nation = dict()
    for _ in k:
        if not k[_] in nation:
            nation[k[_]] = list()
        nation[k[_]].append(_)
	
    # 국가번호만 추출(combinations 모듈을 사용하기 위해서 하는 작업)
    # 다음 작업을 하였을 때 나오는 결과물
    # nation_num:  [-1, 0]
    nation_num = list()    
    for n in nation:
        nation_num.append(n)

    # 국가간 나올 수 있는 조합 구하기
    # 다음 작업을 하였을 때 나오는 결과물
    # nation_combination:  [(-1, 0)]
    nation_combination = list(combinations(nation_num, 2))
   
   
    # 정답 구하기
    answer = 0
    
    for _ in nation_combination:
        answer += len(nation[_[0]]) * len(nation[_[1]])
	
    # nation[-1]에 속해있는 astronaut는 같은 국가가 아니기 때문에 각 astronaut을 다른 국가로 보고 nation[-1]의 조합을 구한다.
    if -1 in nation:
        answer += len(list(combinations(nation[-1], 2)))
        
    return answer

hackos 10개를 사용해서 testcase 두 개를 열람했다.

testcase8을 열람하면 10000 명의 우주 비행사가 존재한다고 나오는데, 분명 답은 제대로 나온다.

그런데 틀렸다고 하는 걸 보니 비효율적으로 짜서 그런 것 같다...

다음에 머리가 좀 더 맑으면 그 때 풀어야겠다..