[그래프] - 팀결성 (서로소 집합 & 크루스칼 알고리즘)

n명의 학생들이 M개의 연산이 주어 진다.

M개의 연산의 종류는 아래와 같다.

1. 팀 합치기

2. 같은팀인지 확인

 

2번과 같은 입력이 들어오면 그에 대한 출력을 Yes or NO로 하라.

 

나의 코드

def union(parent,a,b):
    a=check_parent(parent,a)
    b=check_parent(parent,b)
    if a>b:
        parent[a]=b
    else:
        parent[b]=a 
    
def check_parent(parent,x):
    if parent[x] !=x:
        parent[x]=check_parent(parent,parent[x])
    return parent[x] 

n,m=map(int,input().split())

parent=[0]*(n+1)

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

for i in range(m):
    a,b,c=map(int,input().split())
    if a== 0:
        union(parent,b,c)
    if a==1:
        if  check_parent(parent,b) ==check_parent(parent,c):
            print("YES")
        else:
            print("NO")

전형적인 서로소 집합 알고리즘의 union 함수와 find 함수를 구현하는 문제이다.

N,M의 범위가 모두 100000 까지 가능했으므로 경로 압축 기법을 사용했어야 했다.

 

정답 코드

# 특정 원소가 속한 집합을 찾기
def find_parent(parent, x):
    # 루트 노드가 아니라면, 루트 노드를 찾을 때까지 재귀적으로 호출
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

# 두 원소가 속한 집합을 합치기
def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

n, m = map(int, input().split())
parent = [0] * (n + 1) # 부모 테이블 초기화

# 부모 테이블상에서, 부모를 자기 자신으로 초기화
for i in range(0, n + 1):
    parent[i] = i

# 각 연산을 하나씩 확인
for i in range(m):
    oper, a, b = map(int, input().split())
    # 합치합(Union) 연산인 경우
    if oper == 0:
        union_parent(parent, a, b)
    # 찾기(Find) 연산인 경우
    elif oper == 1:
        if find_parent(parent, a) == find_parent(parent, b):
            print('YES')
        else:
            print('NO')

나의 코드와 완전히 일치 한다 !!