이 문제는 최소신장트리 문제와 동일하다.
크러스컬 알고리즘을 이용한 방법
1. 비용이 적은 순서로 다리를 건설하기 위해 costs를 비용 기준으로 오름차순 정렬한다.
2. 각각의 섬들로 집합을 만들어 connect 리스트에 저장한다. (ex. connect = [{0}, {1}, {2}, {3}...]) 각각의 집합들은 연결되어 있는 섬들의 집합으로 같은 집합에 속해있으면 서로 연결된 관계이다.
3. costs 리스트를 for문을 돌린다. connect에서 현재 연결하려는 두 섬이 포함된 집합을 각각 찾는다. 찾은 두 집합이 같은 집합이면 두 섬은 이미 연결되어있는 상태이므로 다리를 건설하지 않고 continue로 다음 차례로 넘어간다. 두 집합이 다른 집합이면 서로 연결되어 있지 않은 상태이므로 다리를 건설한다. 두 섬이 연결되었음을 업데이트하기 위해 connect에서 아까 찾은 두 집합을 삭제한 후 두 집합의 합집합을 추가한다. answer에는 이 다리를 건설하기 위한 비용을 더한다. 모든 섬이 연결되어 connect의 원소가 하나이고 그 원소인 집합의 길이가 n-1이 되면 break로 for문을 빠져나온다. (모든 섬이 연결되려면 n-1개의 다리가 필요함, while문 쓰는 것도 가능)
def solution(n, costs):
answer = 0
costs.sort(key=lambda x:x[2])
connect = []
for i in range(n):
connect.append({i})
for i in costs:
temp0 = set()
temp1 = set()
for j in connect:
if i[0] in j:
temp0 = j
if i[1] in j:
temp1 = j
if temp0 == temp1:
continue
else:
connect.remove(temp0)
connect.remove(temp1)
connect.append(temp0|temp1)
answer += i[2]
if len(connect)==1 and len(connect[0])==n-1:
break
return answer
처음에 연결된 섬들을 각각의 집합으로 나누지 않고 하나의 집합에 넣었더니 각 섬들의 연결 관계를 알 수 없어 답이 나오지 않았다. 예를 들어, 0-1-2 / 3-4 이렇게 섬들이 연결되어 있을 때 연결 관계에 따라 집합으로 나누면 connect가 [{0,1,2}, {3,4}]가 되지만 하나의 집합만 사용하면 connect가 {0,1,2,3,4}가 되어 모두 연결된 상태와 같게 된다.