union find 응용문제 풀어보면서 분리집합 개념 재활하기

1. 문제

 

16562번: 친구비 (acmicpc.net)

16562번: 친구비

 

2. 풀이

 

어렵게 생각하지 말고.. 먼저 union find 함수를 작성하자

 

먼저 parent를 찾는 find함수

 

def find_parent(x,parent):
    
    if x != parent[x]:
        
        parent[x] = find_parent(parent[x],parent)
    
    return parent[x]

 

 

다음 union을 수행하는 함수

 

def union(a,b,parent):
    
    a = find_parent(a,parent)
    b = find_parent(b,parent)
    
    if parent[a] > parent[b]:
        
        parent[parent[b]] = parent[a]
    
    else:
        
        parent[parent[a]] = parent[b]

 

입력을 받는데.. 학생 번호가 1번부터 시작하니까 친구비에 앞에 [0]을 더해서 인덱스 쉽게 만들어주자

 

n,m,k = map(int,stdin.readline().split())

A = [0] + list(map(int,stdin.readline().split()))

parent = [0]*(n+1)

for i in range(1,n+1):
    
    parent[i] = i

 

그리고 parent를 초기화하고..

 

다음에 친구관계가 주어지는데.. 애초에 주어지는 친구관계끼리는 하나의 집합이고,

 

이들끼리는 친구비가 들지 않으니까

 

이들은 union시켜버려서 하나의 집합을 형성시키자

 

for _ in range(m):
    
    v,w = map(int,stdin.readline().split())

    union(v,w,parent)

 

그러면 친구관계를 이루는 집합이 전부 형성이 되어 있단 말이지

 

준석이가 이 집합들 각각에 친구비를 내서 친구를 사귈려고 시도할거란 말이야

 

근데 각 학생이 요구하는 최소의 친구비만 내면, 집합 내의 모든 학생과 친구가 될 수 있잖아

 

그래서 애초에 union을 시킬때, 대표자를 친구비가 적은 번호가 parent로 들어가도록 만들자

 

def union(a,b,parent):
    
    #a,b의 대표자를 찾는다
    a = find_parent(a,parent)

    b = find_parent(b,parent)
    
    #a의 친구비가 b의 친구비보다 많다면...
    if A[a] > A[b]:
        
        #a가 b에 들어가도록
        #a의 부모를 b로 설정
        parent[a] = b
    
    else:
        
        #a의 친구비가 b의 친구비보다 적다면..
        
        #b가 a에 들어가도록
        #b의 부모를 a로 설정
        parent[b] = a

 

그러면 이제... 각 친구관계 집합의 대표자가 최소 비용을 요구하는 대표로 구성이 되어 있다

 

그래서 parent를 순회해서 각 집합이 요구하는 최소 친구비용을 모두 내서

 

모든 친구와 사귀기를 시도한다

 

모든 친구와 사귈때, 친구비 총합이 가지고 있는 친구비보다 많아지면... 모든 학생을 친구로 만들 수 없고

 

가지고 있는 친구비 내로 할 수 있으면 모든 학생을 친구로 만들 수 있다

 

answer = 0

check = [0]*(n+1)

for p in parent[1:]:
    
    if check[p] == 0:
        
        answer += A[p]
        check[p] = 1

 

0번은 인덱스 접근을 쉽게 하기위해 만든거라 1번부터 순회해서...

 

1번 학생의 대표자부터 접근을 하는거지

 

그 대표자를 사용하지 않았다면... answer에 그 집합과 사귀는 비용을 answer에 누적시키고

 

다음에 중복 누적합 하지 않기 위해 check배열을 1로 바꿔준다

 

근데 이렇게 하면 오답인데

 

여기서 union find를 제대로 이해하지 못했다는거.. 너무 오랜만에 쓴게 티가 난건가

 

--------------------------------------------------------------------------------------------------------------------------

 

union(u,v,parent)를 수행하면...

 

단순히 u와 v만 합쳐지는게 parent에 표시가 되는거지

 

v가 속한 모든 집합의 원소도 u와 합쳐지기는 한데, parent에 표시되는건 별개야

 

예를 들어 

 

친구비가 1 1 1 1 1이고

 

친구 관계 u,v가

 

1 5

2 4

4 3

5 4

 

로 주어진다고 해봐

 

parent의 초기 배열이 [0,1,2,3,4,5]로 초기화 된 상태이고

 

1, 5가 먼저 union되면 parent는

 

[0,1,2,3,4,1]이다

 

다음 2,4가 union되면.. [0,1,2,3,2,1]이다.

 

다음 4,3이 union되면.. 

 

a = 4, b = 3인데... 각각의 대표는 a = 2와  b = 3이다

 

A[2] = A[3]이므로.. parent[3] = 2가 된다.

 

그래서 [0,1,2,2,2,1]이다.

 

다시 5와 4가 union되면... a = 5, b = 4인데 각각의 대표는 a = 1, b = 2이다.

 

A[1] = A[2] 이므로 parent[2] = 1이다.

 

그래서 [0,1,1,2,2,1]이 최종적으로 표시되는 parent이다

 

그런데 union과정을 잘 생각해보면..

 

1 52 44 35 4

 

가 합쳐지면... 1,2,3,4,5가 하나의 집합을 형성하잖아

 

하지만 [0,1,1,2,2,1]을 곧이곧대로 순회한다면... 집합이 1번과 2번 2개있는걸로 보인단 말이지

 

즉 1,2,5와 3,4가 각각 집합으로 된것 같단 말이지

 

하지만 3,4의 parent는 2인데 2의 parent는 1이다.

 

그니까 union함수는 두 원소 u,v만 바꿔서 parent에 표시해주지.. 모든 경우를 다 바꿔주는건 아니다

 

따라서 parent를 순회할때 parent의 대표자를 찾아서 해당 index가 어디 집합에 속하는지까지 찾아야한다

 

answer = 0

check = [0]*(n+1)

for p in parent[1:]:
    
    #p의 대표자를 찾아서 해당 index가 최종적으로 어디 집합에 속하는지를 찾습니다
    p = find_parent(p,parent)
    
    if check[p] == 0:
        
        answer += A[p]
        check[p] = 1

 

그래서 모든 집합에 친구비를 내서 모든 학생을 친구로 만들었을때 드는 비용 answer와 가지고 있는 돈 k를 비교해서..

 

answer이 k보다 크면 모든 학생을 친구로 만들 수 없고 작거나 같으면 모든 학생을 친구로 만들 수 있다

 

from sys import stdin

def find_parent(x,parent):
    
    if x != parent[x]:
        
        parent[x] = find_parent(parent[x],parent)
    
    return parent[x]

def union(a,b,parent):
    
    a = find_parent(a,parent)

    b = find_parent(b,parent)

    if A[a] > A[b]:
        
        parent[a] = b
    
    else:
        
        parent[b] = a

n,m,k = map(int,stdin.readline().split())

A = [0] + list(map(int,stdin.readline().split()))

parent = [0]*(n+1)

for i in range(1,n+1):
    
    parent[i] = i

for _ in range(m):
    
    v,w = map(int,stdin.readline().split())

    union(v,w,parent)

answer = 0

check = [0]*(n+1)

for p in parent[1:]:
    
    p = find_parent(p,parent)
    
    if check[p] == 0:
        
        answer += A[p]
        check[p] = 1

if answer <= k:
    
    print(answer)

else:
    
    print('Oh no')