Journey to the Moon

  • + 0 comments
    from collections import Counter
    
    def journeyToMoon(n, astronaut):
        # Initialize disjoint set (union-find) structures
        root = list(range(n))
        rank = [0] * n
        
        def find_root(a):
            if root[a] != a:
                root[a] = find_root(root[a])  # Path compression
            return root[a]
            
        def union(a, b):
            i, j = find_root(a), find_root(b)
            if i != j:
                if rank[i] > rank[j]:
                    root[j] = i
                elif rank[i] < rank[j]:
                    root[i] = j
                else:
                    root[j] = i
                    rank[i] += 1
        
        # Union all astronaut pairs
        for a, b in astronaut:
            union(a, b)
        
        # Count the size of each connected component
        component_sizes = list(Counter(find_root(i) for i in range(n)).values())
        
        # Calculate the number of valid pairs
        total_pairs = 0
        sum_sizes = 0
        
        for size in component_sizes:
            total_pairs += sum_sizes * size  # Pairs between current and previous components
            sum_sizes += size  # Update the cumulative size sum
        
        return total_pairs