Roads and Libraries

  • + 0 comments
    def roadsAndLibraries2(n, c_lib, c_road, cities):
        # cost 1: c_lib * n
        if c_lib <= c_road:
            return c_lib * n
            
        # cost 2: c_lib * k + c_road * (n - k), where k is num of connected components
        # Initialize parent array for disjoint-set data structure
        parent = list(range(n+1))
        # Initialize rank array for union by rank optimization
        rank = [0] * (n + 1)
        
        # Find operation to recursively find the representative of a set
        def find_parent(city):
            if parent[city] == city:
                return city
            # Path compression: update parent pointers to point directly to the root
            parent[city] = find_parent(parent[city])
            return parent[city]
            
        # Union operation to merge two sets
        def union(city1, city2):
            a, b = find_parent(city1), find_parent(city2)
            if a == b:
                return
            # Union by rank: merge smaller tree into larger tree
            if rank[a] > rank[b]:
                parent[b] = a
            elif rank[a] < rank[b]:
                parent[a] = b
            else:
                parent[b] = a
                rank[a] += 1
        
        # Merge sets based on given roads
        for city1, city2 in cities:
            union(city1, city2)
            
        # Count the number of connected components (sets with representatives)
        k = sum(1 for city in range(1, n+1) if parent[city] == city)
        
        # Calculate cost based on the number of connected components
        return c_lib * k + c_road * (n - k)