Roads and Libraries

  • + 0 comments
    def roadsAndLibraries(n, c_lib, c_road, cities):
        # If library cost is less than or equal to road cost, build a library in each city
        if c_lib <= c_road:
            return c_lib * n
        
        # Initialize disjoint set (union-find) structures
        root = list(range(n + 1))
        rank = [0] * (n + 1)
        
        def find_root(i):
            if root[i] != i:
                root[i] = find_root(root[i])  # Path compression
            return root[i]
        
        def union(a, b):
            i, j = find_root(a), find_root(b)
            if i != j:
                # Union by rank
                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 connected cities
        for a, b in cities:
            union(a, b)
        
        # Count connected components
        k = sum(1 for city in range(1, n + 1) if root[city] == city)
        
        # Calculate total cost: one library per component + roads for remaining cities
        return c_lib * k + c_road * (n - k)