Hackerrank - Journey to the Moon Solution
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.
For example, we have the following data on 2 pairs of astronauts, and 4 astronauts total, numbered through .
1 2
2 3
Astronauts by country are and . There are pairs to choose from: and .
Function Description
Complete the journeyToMoon function in the editor below. It should return an integer that represents the number of valid pairs that can be formed.
journeyToMoon has the following parameter(s):
- n: an integer that denotes the number of astronauts
- astronaut: a 2D array where each element is a element integer array that represents the ID's of two astronauts from the same country
Input Format
The first line contains two integers and , the number of astronauts and the number of pairs.
Each of the next lines contains space-separated integers denoting astronaut ID's of two who share the same nationality.
Constraints
Output Format
An integer that denotes the number of ways to choose a pair of astronauts from different countries.
Sample Input 0
5 3
0 1
2 3
0 4
Sample Output 0
6
Explanation 0
Persons numbered belong to one country, and those numbered belong to another. The UN has ways of choosing a pair:
Sample Input 1
4 1
0 2
Sample Output 1
5
Explanation 1
Persons numbered belong to the same country, but persons and don't share countries with anyone else. The UN has ways of choosing a pair:
Solution in Python
from collections import defaultdict
def group_team(astronaut):
country = defaultdict(set)
for i,(a,b) in enumerate(astronaut):
country[i].update([a,b])
while country:
i = list(country.keys())[0]
x = country[i]
pops = [k for k,v in country.items() if v & x]
if len(pops) ==1:
country.pop(i)
yield x
else:
for p in pops:
x = x | country.pop(p)
country[i] = x
def journeyToMoon(n, astronaut):
arr = [len(i) for i in group_team(astronaut)]
s = sum(arr)
c = 0
for i in range(s,n):
arr.append(1)
s+=1
for i in arr:
s-=i
c+= i*s
return c
n,p = map(int,input().split())
astronaut = [list(map(int,input().split())) for _ in range(p)]
print(journeyToMoon(n, astronaut))
Answer explanation
First we create a helper function group_team which associates astronauts to same group if they have any matching values.
For example [0,1] , [2,3], [0,4] will return
[{0,1,4}, {2,3}]
Now in next step we will convert this list values to count of items inside the values and our list will become.
[3,2]
Let's assume n = 8
As we can see that we have only total of 5 astronauts instead of 8 we will assume the 3 astronauts which isn't included in our list as a team of individuals.
Now our list of teams will have the following count of astronauts in each team.
[3,2,1,1,1]
Now we just have to move from left to right.
Team 1 has 3 astronauts, and each of those astronauts can be paired with astronauts in other teams [2,1,1,1]
There can be 3*(2+1+1+1) = 3*5 = 15 different ways to do so
We move to Team 2.
Team 2 has 2 astronauts, and each of those astronauts can be paired with astronauts in other teams [1,1,1]
There can be 2*(1+1+1) = 2*3= 6 different ways to do so
Since we have already calculated all the possible pairings with team 1. We will no longer include team 1(items in the left).
We move to Team 3.
Team 3 has 1 astronauts, and each of those astronauts can be paired with astronauts in other teams [1,1]
There can be 1*(1+1) = 1*2= 2 different ways to do so
Team 4 and Team 5 can together make a pair of 1.
Now we just have to add all these values and we will get the total number of possibilities 15+6+2+1= 24