HackerRank Solution: Journey to the Moon

moon_journey.py

#!/bin/python3
# https://www.hackerrank.com/challenges/journey-to-the-moon/problem

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

# Disjoint sets : search & flattening
def get_root(node, sets):
    path = [node]
    while type(sets[path[-1]]) is int:
        path.append(sets[path[-1]])
    # disjoint set : collapse tree
    for node in path[:-1]:
        sets[node] = path[-1]
    return path[-1]


def journeyToMoon(n, astronaut):
    node_to_set = {node: {node} for node in range(n)}

    for a, b in astronaut:
        a_root, b_root = get_root(a, node_to_set), get_root(b, node_to_set)
        if a_root == b_root:
            continue
        else:
            a_set, b_set = node_to_set[a_root], node_to_set[b_root]
            a_set.update(b_set)
            for member in b_set:
                node_to_set[member] = a_root

    # Rather than just computing all combinations, this is optimized in case of 
    # a low number of edges, high number of nodes
    
    disjoint_sets = [ds for node, ds in node_to_set.items() if type(ds) is set]
    
    # Partition into single-node sets and multi-node sets
    singleton_sets = [ds for ds in disjoint_sets if len(ds) == 1]
    multi_node_sets = [ds for ds in disjoint_sets if len(ds) > 1]
    
    multi_set_pairs = combinations(multi_node_sets, 2)
    
    singleton_set_pair_count = (len(singleton_sets) * (len(singleton_sets) - 1) // 2)
    return singleton_set_pair_count + \
        sum(len(x) * len(y) for x, y in multi_set_pairs) + \
        sum(len(ds) * (len(disjoint_sets) - len(multi_node_sets)) for ds in multi_node_sets)


# ------ HackerRank Starter Code ------
if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    np = input().split()

    n = int(np[0])

    p = int(np[1])

    astronaut = []

    for _ in range(p):
        astronaut.append(list(map(int, input().rstrip().split())))

    result = journeyToMoon(n, astronaut)

    fptr.write(str(result) + '\n')

    fptr.close()