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 Counter
# A class to represent a disjoint set

class DisjointSet:
parent = {}

# perform MakeSet operation
def makeSet(self, universe):

# create n disjoint sets (one for each item)
for i in universe:
self.parent[i] = i

# Find the root of the set in which element k belongs
def Find(self, k):

# if k is root
if self.parent[k] == k:
return k
# recur for the parent until we find the root
return self.Find(self.parent[k])

# Perform Union of two subsets
def Union(self, a, b):

# find the root of the sets in which elements
# x and y belongs
x = self.Find(a)
y = self.Find(b)

self.parent[x] = y

def getSets(universe, ds):
return [ds.Find(i) for i in universe]

def journeyToMoon(no_of_astronauts):

# universe of items i.e astronauts
universe = list(range(no_of_astronauts))

# initialize disjoint set
ds = DisjointSet()

# create a singleton set for each astronaut
ds.makeSet(universe)

for pair in pairs:
# Union all astronaut that are from the same country
ds.Union(*pair)

sets = getSets(universe, ds)

# Count the no. of astronauts in each country
astronauts_each_country = Counter(sets)

# initialize no. of valid pairs i.e pair of 2 astronauts from different countries
no_of_valid_pairs = 0

remaining_astronauts = no_of_astronauts
# Count possible pairs
for astronauts_in_country in astronauts_each_country.values():
remaining_astronauts -= astronauts_in_country
no_of_valid_pairs += astronauts_in_country * remaining_astronauts

return no_of_valid_pairs

# DisjointSet data structure (UnionFind algorithm)
if __name__ == '__main__':

# Take input from user
no_of_astronauts, no_of_pairs = map(int, input().split())
pairs = []
for i in range(no_of_pairs):
pair = list(map(int, input().split()))
pairs.append(pair)

print(journeyToMoon(no_of_astronauts))