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))