[백준 1647번] 도시 분할 계획- python / 크루스칼 알고리즘

문제 링크

https://www.acmicpc.net/problem/1647

1647번: 도시 분할 계획

 

풀이 

- 크루스칼 알고리즘을 사용해서 모든 정점을 잇는 최소 비용을 구한 후 가장 큰 값을 빼줘서 두 도시로 나눌 수 있다. 

set을 사용해서 구현하니 시간 초과가 발생하였다. 

# 시간 초과 발생 
import sys
input = sys.stdin.readline 

n, m = map(int,input().split())

cost = []

for _ in range(m):
    a, b, c = map(int,input().split())
    cost.append([a,b,c])
    
cost.sort(key = lambda x : x[2])
connect = set([cost[0][0], cost[0][1]])
answer = []
answer.append(cost[0][2])

while len(connect) < n :
    for c in cost:
        if c[0] in connect and c[1] in connect :
            continue 
            
        if c[0] in connect or c[1] in connect :
            connect.update([c[0],c[1]])
            answer.append(c[2])
            break
        
print(sum(answer[:-1]))

- 유니온 파인드 함수로 변경해서 구현하니 통과 

# https://www.acmicpc.net/problem/1647
import sys
input = sys.stdin.readline 

# 부모 찾는 함수 
def find(x):
    if x == parent[x]:
        return x
    parent[x] = find(parent[x])
    
    return parent[x]

# 두 집합 합치는 함수 
def union(a,b):
    a = parent[a]
    b = parent[b]
    
    if a == b:
        return
    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        
n, m = map(int,input().split())

# 부모 저장 테이블 
parent = [0] * (n+1)

# 자기 자신으로 초기화 
for i in range(1,n+1):
    parent[i] = i 
    
edges = []
for _ in range(m):
    a, b, cost = map(int,input().split())
    edges.append([cost,a,b])

# cost 기준 정렬 
edges.sort()
# 연결 비용 저장
answer = []

for edge in edges:
    cost, a, b = edge
    
    # 같은 집합에 속하지 않으면 
    if find(a) != find(b):
    	# 연결하고 비용 더하기 
        union(a,b)
        answer.append(cost)
    
 # 가장 큰 비용 제거 
print(sum(answer[:-1]))