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

문제 내용.

 

유니온 파인드 문제.

각 물탱크 집합에 들어있는 청정수 개수와 고인물 개수를 모두 세어두고,

두 물탱크를 연결할 때 해당 물탱크가 서로 다른 집합에 있다면 두 물탱크 집합의 청정수 개수와 고인물 개수를 각각 합친다.

그 후에는 쿼리가 들어올 때마다 물탱크 집합을 찾아준 뒤 청정수 개수와 고인물 개수를 비교해 출력하면 된다.

 

코드는 다음과 같다.

 

# -*- coding: utf-8 -*-
import sys
from collections import deque
import heapq
import bisect
import math
from itertools import product
"""
from itertools import combinations
from itertools import combinations_with_replacement
from itertools import permutations
import copy
"""
#input=sys.stdin.readline
#print=sys.stdout.write
#sys.setrecursionlimit(100000000)
n,m,q=map(int,input().split())
L=list(map(int,input().split()))
parent=[i for i in range(n+1)]
count=[[0,0] for i in range(n+1)]

for i in range(1,n+1):
    count[i][L[i-1]]+=1

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

def union(x,y):
    x=find(x)
    y=find(y)
    x,y=min(x,y),max(x,y)
    if x!=y:
        parent[y]=x
        count[x][0]+=count[y][0]
        count[x][1]+=count[y][1]
        
for i in range(m):
    u,v=map(int,input().split())
    union(u,v)
    
for i in range(q):
    a=find(int(input()))
    if count[a][0]<count[a][1]:
        print(1)
    else:
        print(0)

뭔가 학회 채점 현황에서 자주 본 것 같았는데, 풀고 보니 우리 학교 청정수컵 문제였다.