[프로그래머스] 섬 연결하기 LV3[prim, union-find]

https://school.programmers.co.kr/learn/courses/30/lessons/42861

프로그래머스

 

# Prim 알고리즘

heapq 사용x, 정통알고리즘

INF = float('INF')
def prim(n, costs, x):
    selected = [False]*n # 우리편
    distance = [INF]*n # 가중치를 담는 배열  
    adj_mat = [[INF]*n for _ in range(n)] # 인접행렬
    
    # 1)graph init 
    for c in costs:
        adj_mat[c[0]][c[1]] = adj_mat[c[1]][c[0]] = c[2] # 가중치
    distance[x] = 0 # 초기화    
    # 2) 비즈니스 로직
    for i in range(n): 
        # distance의 최솟값 찾기
        minnum = [0, INF]
        for idx, value in enumerate(distance):
            if not selected[idx] and value < minnum[1]:
                minnum[0] = idx # 인덱스
                minnum[1] = value # 값
                
        selected[minnum[0]] = True # 우리편으로
        # if distance[minnum[0]] == INF: return
        for idx, value in enumerate(adj_mat[minnum[0]]):
            if adj_mat[minnum[0]][idx] != INF:
                if not selected[idx] and value < distance[idx]: 
                    distance[idx] = value # 최솟값으로 업데이트 
  
    return sum(distance)

def solution(n, costs):
    result = prim(n, costs, 0) # 아무데서나 시작, 모든 가중치 더하면, 가장 적은 비용으로 통행하는 가중치
    
    return result

---

# Prim 알고리즘

heapq 사용O, 인접리스트 이용.

# 3) prim, heapq(O)
# 인접리스트 이용
from heapq import heappush, heappop
def solution(n, costs):
    adj_list = [[] for _ in range(n)]
    for cost in costs:
        adj_list[cost[0]].append((cost[1], cost[2]))
        adj_list[cost[1]].append((cost[0], cost[2]))
    
    start = 0
    visited = [False] * n
    visited[start] = True
    heap = []
    result = 0
    
    for node, weight in adj_list[start]:
        heappush(heap, (weight, start, node))
    
    while heap:
        weight, start, node = heappop(heap)
        if visited[node]:
            continue
        
        visited[node] = True
        result += weight
        
        for next_node, next_weight in adj_list[node]:
            heappush(heap, (next_weight, node, next_node))
        

    return result

---

# 크루스칼(union-find) 알고리즘

# 1) 크루스칼
# 이 풀이는 이중for문이 되지만 매우 빠른듯? 코드도 간단.

# 1.정렬하기
# 2. 이미 낮은 걸로 연결되어 있으면 무시
# 3. 연결 안되어 있으면 연결
# => find연산이 in임, 같은집합인지 물어보는거임.

def solution(n, costs):
    answer = 0
    costs.sort(key = lambda x: x[2])
    link = set([costs[0][0]]) # 시작연결점
    
    # 크루스칼로 최소 비용 구하기
    while len(link) != n:
        for v in costs:
            if v[0] in link and v[1] in link: 
                continue
            if v[0] in link or v[1] in link:  # 1개가 연결안되어있으면 
                link.update([v[0], v[1]]) # 중복제거해서 들어감.
                answer += v[2] # 가중치 
                break 
    return answer