[Python/BOJ] 백준 17398 통신망 분할_유니온파인드

코테에 유니온파인드가 나오나...?라는 의구심을 품고 푼 문제 

 

유니온 파인드 문제를 처음 풀어봐서 주석을 많이 달아놨다. 

 

 

이 문제를 푸는 방법은

끊을 간선을 제외한 모든 간선을 먼저 연결하고 끊을 간선들을 역순으로 하나씩 연결한다.

 

배열의 부모노드의 값은 음수로 집합의 개수가 저장되어있다는 것이 중요하다. 

 

 

 

---

 

 

<Main> 

처음 배열은 -1로 초기화한다. 

 

edges,div로 연결할 리스트와 제거할 리스트를 입력받는다.

연결할 리스트의 개수만큼을 돌면서 제거될 연결이 아니라면 연결해준다.(union)

 

answer 변수를 0으로 초기화한다.

제거할 연결을 역순으로 union한다. union은 각 집합의 개수의 곱을 반환하므로 반환값을 answer에 더한다.

 

answer을 출력한다. 

 

 

 

<union>

; 각 원소를 연결해주고 두 그룹의 크기의 곱을 반환해주는 함수 

 

find를 통해서 부모노드의 번호를 구한다.

부모노드의 값은 음수로 집합의 개수를 갖고있으므로 두개의 값을 곱하고 res로 저장해둔다.

(둘다 음수이므로 곱하면 양수가 되어서 우리가 원하는 값이 잘 나온다.)

 

만약 부모노드가 동일하다면 두 원소는 같은 집합의 원소이므로 return 0을 한다. 

만약 부모노드가 동일하지 않다면 한 부모노드의 값을 다른 부모 노도의 값에 더하여 저장한다. 

각 집합의 원소의 개수를 저장하고 있기에 더하는 것이다. 이제 두 집합은 하나의 집합이다. 

그리고 더함을 받지않은 부모노드는 이제 더함을 받은 부모노드의 번호를 값으로 갖고있다. 

 

 

 

<find>

; 부모 노드의 번호를 반환하는 함수 

부모노드라면 집합의 개수가 음수로 저장되어있다는 것을 생각하자.

만약 값이 음수라면 부모노드라는 뜻이므로 번호를 반환한다.

만약 값이 음수가 아니라면 부모노드가 아닌 노드라는 의미이므로 

재귀를 통해서 부모노드를 찾고 그 과정에서 부모노드의 값으로 배열의 값이 변경된다. 

 

 

1번은 -4로 4개의 원소를 가진 집합의 부모노드라는 것을 알 수 있다. 

2,3,4번은 1번이 부모노드인 집합의 구성원임을 알 수 있다. 

 

 

 

코드 : 

# u배열, a와 b를 union
def union(u, a, b):
    
    # 각각의 부모를 구함 
    a = find(u, a)
    b = find(u, b)

    # union할 두 그룹의 크기의 곱
    res = u[a] * u[b]
    # 같은 집합이라면 
    if a == b :
        return 0
    # 같은 집합이 아니라면 
    else :
        u[a]+=u[b]
        u[b] = a 
    return res


# 부모 노드를 찾는 함수 
def find(u, x):
    # 부모노드라면 
    if u[x] < 0:
        return x
    # 부모노드가 아니라면
    # 재귀를 이용해서 부모노드를 찾음 
    u[x] = find(u, u[x])
    return u[x]

if __name__=="__main__" :
    n, m, q = map(int, input().split())
    u = [-1] * (n + 1)
    
    # 끊을 간선을 제외한 모든 간선을 먼저 연걸해주고 끊을 간선들을 
    # 역순으로 하나씩 연결 
    # 부모노드에 음수로 갯수를 저장한다. 
    
    # 연결할 리스트 
    edges = [list(map(int, input().split())) for _ in range(m)]
    # 제거할 리스트 
    div = [int(input()) for _ in range(q)]
    div_set = set(div)


    # 제거되지 않는 연결들에 대해 union 처리
    for i in range(m):
        # 제거할 간선이라면 pass 
        if (i+1) in div_set:
            continue
        # 제거할 간선이 아니라면 연결 
        union(u, edges[i][0], edges[i][1])

    # 연결을 끊는데 들어가는 비용
    answer = 0

    # 제거할 연결을 역순으로 union
    for i in range(q-1, -1, -1):
        answer += union(u, edges[div[i]-1][0], edges[div[i]-1][1])

    print(answer)

 

 

 

 

참고 영상 : 

https://www.youtube.com/watch?v=AMByrd53PHM&t=107s