Journey to the Moon

  • + 0 comments

    Python 3

    from collections import Counter
    
    def journeyToMoon(n: int, astronaut: list[tuple[int, int]]) -> int:
        parent = list(range(n))
    
        def find(i):
            if parent[i] == i:
                return i
            parent[i] = find(parent[i])
            return parent[i]
    
        def union(i, j):
            root_i = find(i)
            root_j = find(j)
            if root_i != root_j:
                parent[root_i] = root_j
        for a, b in astronaut:
            union(a, b)
        country_sizes = Counter(find(i) for i in range(n)).values()
        total_pairs = (n * (n - 1)) // 2  
        same_country_pairs = sum(size * (size - 1) // 2 for size in country_sizes)
    
        return total_pairs - same_country_pairs