[이코테] 그래프 이론_팀 결성 (python)

나의 풀이

1. 서로소 집합 알고리즘(union-find 알고리즘) 풀이

# union 연산(팀 합치기) :두 학생이 속한 팀을 합치기
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
        
# find 연산: 특정 학생이 어떤 팀에 속하는지 찾기
def find_parent(parent, x):
	# 루트 노드가 아니라면, 루트 노드를 찾을 때까지 재귀적으로 호출
	if parent[x] != x:
		parent[x] = find_parent(parent, parent[x])
        
	return parent[x]
    
# 같은 팀 여부 확인 연산: 두 학생이 같은 팀에 속하는지 확인
def is_same_parent(a, b):
	# 두 학생이 같은 팀에 속할 경우
	if find_parent(parent, a) == find_parent(parent, b):
		return print('YES')
        
	else: # 두 학생이 서로 다른 팀에 속할 경우
		return print('NO')
        
n, m = map(int, input().split()) # n: 0 ~ n번까지 학생, m: 연산의 수
parent = list(range(n + 1)) # 0 ~ n번까지 부모를 자기 자신으로 초기화
    
for i in range(m):
	# operation: 연산의 종류, a 학생, b 학생
	operation, a, b = map(int, input().split())
    
	if operation == 0: # union 연산인 경우
		union_parent(parent, a, b)
        
	else: # 같은 팀 여부 확인 연산인 경우
		is_same_parent(a, b)

• 부모 테이블에서, 자기 자신을 부모로 초기화 하는 경우..
  기존에는

parent = [0] * (n + 1) # 부모 테이블 초기화

for i in range(n + 1): # 0 ~ n번까지 부모를 자기 자신으로 초기화
	parent[i] = i

로 했지만

parent = list(range(n + 1))

로 하면 된다!