Roads and Libraries

  • + 0 comments
    from collections import defaultdict, deque
    
    def roadsAndLibraries(n, c_lib, c_road, cities):
        """
        Cost 1 : one lib in each city
        Formula: c_lib * n
        
        Cost 2 : one lib per group of connected cities + connecting roads
        Formula: c_lib * k + c_road * (n - k), where k is number of groups
        
        If c_lib <= c_road, Cost 1 is lower; otherwise, Cost 2 is.
        """
        # Return Cost 1
        if c_lib <= c_road:
            return c_lib * n
    
    
        # Calculate Cost 2 below:
        
        # Initialize number of groups
        k = 0
        
        # Record direct connections for each city
        connections = defaultdict(set)
        for a, b in cities:
            connections[a].add(b)
            connections[b].add(a)
        
        # Track visited cities
        visited = [False] * (n + 1)
        
        # Identify connected components (groups of cities)
        for city in range(1, n + 1):
            if visited[city]:
                continue
            
            # Start a new group
            visited[city] = True
            k += 1
            
            # BFS to mark all cities in this group
            dq = deque([city])
            while dq:
                curr_city = dq.popleft()
                for linked_city in connections[curr_city]:
                    if not visited[linked_city]:
                        visited[linked_city] = True
                        dq.append(linked_city)
                        
        # Return Cost 2
        return c_lib * k + c_road * (n - k)