백준 16562 : 친구비 (Python)

본문 링크

📌 내 풀이

import sys
input=sys.stdin.readline
sys.setrecursionlimit(10**9)

def Find(x):

    if disjoint[x]!=x:
        disjoint[x]=Find(disjoint[x])
    return disjoint[x]

def Union(a,b):

    a=Find(a)
    b=Find(b)

    if a>b:
        disjoint[a]=b
    else:
        disjoint[b]=a

N,M,K=map(int,input().split())
Money=list(map(int,input().split()))

disjoint=[0]*(N+1)
visit=[False]*(N+1) # 모든 친구를 친구로 만들어야 함.

for i in range(1,N+1):
    disjoint[i]=i

People=[]
for i in range(M):

    a,b=map(int,input().split())

    Union(a,b)
    People.append([a,b])

for i in range(M):
    Union(People[i][0] , People[i][1])

check=[ [] for _ in range(N+1) ]

for i in range(1,N+1):
    check[disjoint[i]].append(Money[i-1])

total=0

for i in range(1,N+1):
    if check[i]:
        total+=min(check[i])

if total<=K:
    print(total)
else:
    print("Oh no")

📌 수정한 풀이

import sys
input=sys.stdin.readline
sys.setrecursionlimit(10**9)

def Find(x):

    if disjoint[x]!=x:
        disjoint[x]=Find(disjoint[x])
    return disjoint[x]

def Union(a,b):

    a=Find(a)
    b=Find(b)

    if Money[a]<Money[b]:
        disjoint[b]=a
    else:
        disjoint[a]=b

N,M,K=map(int,input().split())
Money=[0]+list(map(int,input().split()))

disjoint=[0]*(N+1)

for i in range(1,N+1):
    disjoint[i]=i

for i in range(M):

    a,b=map(int,input().split())

    Union(a,b)
total=0

tmp = set()
for i in range(1,N+1):
    if Find(i) not in tmp:
        total+=Money[disjoint[i]]
        tmp.add(disjoint[i])

if total<=K:
    print(total)
else:
    print("Oh no")

📌 어떻게 접근할 것인가?

일단 일반적인 유니온 파인드 코드로 구성하면 됩니다. 집합을 합치고 disjoint 배열 값에 따라서 친구비용의 최소값을 구해주었습니다.

다만 유니온 파인으로 집합을 합치는 과정에서 문제가 생겼습니다.

• 5 4 100
  1 1 1 1 1
  1 5
  2 4
  4 3
  5 4

이런 입력이 주어졌을때

실제로 disjoint 배열은 이렇게 됩니다. 즉 완전히 합쳐지지 않은 것이죠

저는 이 문제를 유니온을 2번을 해서 풀었지만 총 2가지 과정을 통해 해결할수 있습니다.

• 최소값을 구하기 때문에 유니온할때 노드번호가 작은값을 넣지말고 값이 작은값을 우선시 해서 넣어준다.

tmp = set()
for i in range(1,N+1):
    if Find(i) not in tmp:
        total+=Money[disjoint[i]]
        tmp.add(disjoint[i])

• $tmp$ 를 만들어 주어서 방문하지 않은 지점에 대해 최소값을 더해준다.