[hackerrank][python]Roads and Libraries

https://www.hackerrank.com/challenges/torque-and-development/problem

#!/bin/python3
from pprint import pprint
import sys
from collections import defaultdict

def dfs(graph, start):
  visited, stack = set(), [start]
  while stack:  # dfs는 stack을 이용한다.
    vertex = stack.pop()
    if vertex not in visited:
      visited.add(vertex)
      stack.extend(graph[vertex] - visited)  # 해당 정점에 연결된 것들을 stack에 넣는다.
  return visited
  

def roadsAndLibraries(n, c_lib, c_road, cities):
  total_cost = 0
  
  if c_road >= c_lib or len(cities) == 0: 
    return n*c_lib  # 도서관이 싸니까 도로를 만들 이유가 없다.
  else:  # here be real shits
    city_graph = defaultdict(set)  # 저절로 set이 만들어지기 때문에 add만 하면 됨.
    connected_cities = set() # might make this into a set
    for i in cities:
      city_graph[i[0]].add(i[1])
      city_graph[i[1]].add(i[0])
      connected_cities.add(i[0])
      connected_cities.add(i[1])
    
    # print(city_graph)
    # print(connected_cities)
    isolated_cities = set([i+1 for i in range(n)]) - connected_cities
    
    while connected_cities:  # dfs를 이용하여 도시들을 돌아다닌다. 그리고 도시들이 비어질 때까지 돌아다닌다. 그리고 도시들은 한 묶음(여러그룹일 것이다.)에 끝나지 않는 경우가 많기 때문에 dfs를 여러번 실행해야 한다.
    # if connected_cities:  # it's wrong. Becuase the cities will have multiple groups.
      temp = connected_cities.pop()
      component = dfs(city_graph, temp)
      # print('component', component)
      total_cost += (len(component) - 1)*c_road + c_lib  # 라이브러리 하나에 도로 여러개
      connected_cities -= component  # 지들끼리  연결된 component connected_ciites에서 제외
    return total_cost + len(isolated_cities)*c_lib
    

# 모든 도시가 도서관을 접할 수 있도록 해야 한다.
# 모든 도시에 도서관을 만들거나
# 도서관을 하나 만들어서 연결되어 있는 도시(정점)끼리 도로를 만들어야 한다.
# (도로가 연결되어 있지 않다면 도서관을 연결해야 하겠지...)
#if __name__ == "__main__":
q = int(input().strip())
for a0 in range(q):
    n, m, c_lib, c_road = input().strip().split(' ')
    n, m, c_lib, c_road = [int(n), int(m), int(c_lib), int(c_road)]
    cities = []
    for cities_i in range(m):
       cities_t = [int(cities_temp) for cities_temp in input().strip().split(' ')]
       cities.append(cities_t)

    result = roadsAndLibraries(n, c_lib, c_road, cities)
    print(result)