문제
- 학생이 N명인 학교에서 K원을 가진 준석이가 가장 적은 비용으로 모두와 친구가 되는 방법은?
- 1 <= N <= 10,000
- 1 <= K <= 10,000,000
- 학생 i 에게 $A_i$ 만큼의 돈을 주면 그 학생과 1달간 친구가 된다.
- 친구의 친구도 친구다. (친구 관계 수는 M개 있다)
- 준석이가 친구를 다 사귈 수 없다면 "Oh no" 출력한다.
풀이 1
- 비용이 적은 학생부터 시작 + 그래프 탐색으로 경우의 수 줄이기
- => 연결된 친구 관계에서 가장 비용이 적은 친구에게만 돈을 주면 된다.
- 시간 복잡도 : O(N * (N+E)) $\simeq$ O($N^2$) => 가능하다!
import heapq
import sys
sys.setrecursionlimit(10**9) # 재귀 제한
# 친구 관계 확인하기
def check(graph, x, visited):
for i in graph[x]:
if not visited[i] :
visited[i] = True
check(graph, i, visited)
N, M, K = map(int, input().split()) # 학생 수, 친구 관계 수, 돈
price = list(map(int, input().split())) # 친구비
# 친구 관계 그래프 - 중복 제거를 위해 set으로
graph = [ set() for _ in range(N+1) ]
for _ in range(M):
a, b = map(int, input().split())
graph[a].add(b)
graph[b].add(a)
# 친구비가 적은 학생 순으로 정렬하기
price_sort = []
for i in range(N):
heapq.heappush(price_sort, (price[i], i+1))
# 한 명씩 친구 사귀기
money = 0
visited = [False] * (N+1)
while price_sort:
cost, now = heapq.heappop(price_sort)
if not visited[now]: # 아직 친구가 아니라면
money += cost # 비용 지불
visited[now] = True # 방문 표시
check(graph, now, visited) # 친구의 친구는 내 친구~
print(money) if money <= K else print("Oh no")
풀이 2
- 각 친구 관계에 대한 그래프를 만든다는 점에서 유니온 파인드로 해결할 수 있다.
- 친구의 친구 관계로 서로 연결된 그래프를 만들고 그 중 가장 최소 비용을 가진 친구를 대장(부모)으로 둔다.
# 친구 대장 찾기
def find_parent(parent, x):
if parent[x] != x:
parent[x] = find_parent(parent, parent[x])
return parent[x]
# 친구 집합 묶기
def union_parent(parent, a, b, price):
a = find_parent(parent, a)
b = find_parent(parent, b)
if price[a] < price[b] :
parent[b] = a
else:
parent[a] = b
# 입력 받기
N, M, K = map(int, input().split()) # 학생 수, 친구 관계 수, 돈
price = [0] + list(map(int, input().split())) # 친구비
parent = [i for i in range(N+1)] # 각 친구 그룹의 대장
# 친구 그룹 별로 묶기
for _ in range(M):
a, b = map(int, input().split())
union_parent(parent, a, b, price)
money = 0
friend = set()
for i in range(1, N+1):
x = find_parent(parent, i) # 속한 친구 그룹의 대장찾기
if x not in friend : # 아직 친구가 아니면
money += price[x]
friend.add(x)
print(money) if money <= K else print("Oh no")
References
