Roads and Libraries

  • + 0 comments
    def roadsAndLibraries(n, c_lib, c_road, cities):
        # if libraries are cheaper we dont need to build roads
        if c_lib < c_road:
            return c_lib * n
        visited = [0] * n
        transport_stack = []
        cities_map = defaultdict(set)
        islands = 0
        # create a graph of road destinations
        # add both ways for bidirectional paths
        for i in cities:
            cities_map[i[0]-1].add(i[1]-1)
            cities_map[i[1]-1].add(i[0]-1)
        # iterate over the list of road edges
        for city in cities_map:
            if visited[city]:
                continue
            # if we have never been at this city before it is a new island
            islands += 1
            transport_stack.append(city)
            # travel the path of the road
            # if we find a new city add it to the stack
            while transport_stack:
                new_destination = transport_stack.pop(-1)
                if not visited[new_destination]:
                    visited[new_destination] = 1
                    if new_destination in cities_map:
                        for destination in cities_map[new_destination]:
                            if not visited[destination]:
                                transport_stack.append(destination)
    		# if we nevered visited after looking at all the roads then it must be an island 
        islands += n - sum(visited)
        roads = n - islands
        # we need 1 road per city and 1 library per island
        result = (roads * c_road) + (c_lib * islands)
        return result