백준 25187 : 고인물이 싫어요 (Python)

본문 링크

"""
같은 두 물탱크를 연결하는 파이프가 여러 개 존재할 수 있다.
- 중복처리해야함.

"""
import sys
input=sys.stdin.readline


def Find(x):

    if x!=disjoint[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
        dp[b]+=dp[a]
    elif a<b:
        disjoint[b]=a
        dp[a]+=dp[b]

N,M,Q=map(int,input().split())

disjoint=[0]*(N+1)

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

L=list(map(int,input().split()))
dp=[0]*(N+1)

for i in range(1,N+1):
    if L[i-1]==1:
        dp[i]=1
    else:
        dp[i]=-1


for i in range(M):
    a,b=map(int,input().split())

    Union(a,b)

for i in range(Q):

    K=int(input())

    if dp[Find(K)]>=1:
        print(1)
    else:
        print(0)

📌 어떻게 접근할 것인가?

문제에서 청정수를 얻을 수 있는지를 판별하는 것입니다.

• 청정수가 담긴 물탱크의 수가 고인물이 담긴 물탱크의 수보다 더 많은 경우 청정수를 얻을 수 있다.

유니온 파인드를 사용하여서 풀었습니다. 문제에서 청정수라면 $1$ , 고인물이라면 $0$ 을 주는데
유니온 파인드를 사용해서 배열의 값을 합쳐야 하기 때문에 고인물을 $-1$ 로 고쳤습니다.

이후 $Union$ 함수를 다음과 같이 구현했습니다.

def Union(a,b):

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

    if a>b:
        disjoint[a]=b
        dp[b]+=dp[a]
    elif a<b:
        disjoint[b]=a
        dp[a]+=dp[b]

다만 중복되는 수가 들어올 수 있기 때문에 a==b 일때는 넘겨야 합니다.

이후 $dp$ 배열의 값이 1보다 크거나 같으면 청정수가 더 많으므로 청정수를 얻을 수 있습니다.

이는 다음과 같이 구현했습니다

    if dp[Find(K)]>=1:
        print(1)
    else:
        print(0)

$Find(K)$ 를 한 이유는 현재 $dp$ 에서는 번호가 작은 순서에 더 큰 값을 넣었기 때문에 즉 가장 번호가 작은 부모노드에 값을 저장을 했기 때문에 $dp[K]$ 가 아니라 $dp[Find(K)]$ 사용해야합니다.