##### Short Problem Definition:

The member states of the UN are planning to send *2* 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.

##### Link

##### Complexity:

time complexity is O(a(N))

space complexity is O(N)

##### Execution:

The algorithm has two parts. Part one calculates the size of all countries with the use of Union-Find. Both operations on the data structure take O log-star time. The second part derives the number of all pairs based on country sizes. That algorithm is very similar to Handshakes.

I used the UF implementation from ActiveState. For more cool uses of Union-Find, see WireBurnouts.

##### Solution:

#!/bin/python import math import os import random import re import sys from collections import defaultdict class UnionFind: '''http://code.activestate.com/recipes/215912-union-find-data-structure/''' def __init__(self): self.num_weights = {} self.parent_pointers = {} self.num_to_objects = {} self.objects_to_num = {} self.__repr__ = self.__str__ def insert_objects(self, objects): for object in objects: self.find(object); def find(self, object): if not object in self.objects_to_num: obj_num = len(self.objects_to_num) self.num_weights[obj_num] = 1 self.objects_to_num[object] = obj_num self.num_to_objects[obj_num] = object self.parent_pointers[obj_num] = obj_num return object stk = [self.objects_to_num[object]] par = self.parent_pointers[stk[-1]] while par != stk[-1]: stk.append(par) par = self.parent_pointers[par] for i in stk: self.parent_pointers[i] = par return self.num_to_objects[par] def union(self, object1, object2): o1p = self.find(object1) o2p = self.find(object2) if o1p != o2p: on1 = self.objects_to_num[o1p] on2 = self.objects_to_num[o2p] w1 = self.num_weights[on1] w2 = self.num_weights[on2] if w1 < w2: o1p, o2p, on1, on2, w1, w2 = o2p, o1p, on2, on1, w2, w1 self.num_weights[on1] = w1+w2 del self.num_weights[on2] self.parent_pointers[on2] = on1 # Complete the journeyToMoon function below. def journeyToMoon(n, astronauts): # generate disjoint sets uf = UnionFind() uf.insert_objects(xrange(n)) for astronaut in astronauts: uf.union(astronaut[0], astronaut[1]) # calculate sizes of countries country_sizes = defaultdict(int) for i in xrange(n): country_sizes[uf.find(i)] += 1 # calculate number of viable pairs sm = 0 result = 0 for size in country_sizes.values(): result += sm*size; sm += size; return result if __name__ == '__main__': fptr = open(os.environ['OUTPUT_PATH'], 'w') np = raw_input().split() n = int(np[0]) p = int(np[1]) astronaut = [] for _ in xrange(p): astronaut.append(map(int, raw_input().rstrip().split())) result = journeyToMoon(n, astronaut) fptr.write(str(result) + '\n') fptr.close()

If you enjoyed this post, then make sure you subscribe to my Newsletter and/or Feed.