[이것이 코딩 테스트다] with Python
[이것이 코딩 테스트다] 팀 결성 with Python
Evolving Developer
2023. 6. 5. 19:12
🔒 문제
학교에서 학생들에게 0번부터 N번까지의 번호를 부여했다. 처음에는 모든 학생이 서로 다른 팀으로 구분되어, 총 N+1개의 팀이 존재한다. 이때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.
1. '팀 합치기' 연산은 두 팀을 합치는 연산이다.
2. '같은 팀 여부 확인' 연산은 특정한 두 학생이 같은 팀에 속하는지를 확인하는 연산이다.
선생님이 M개의 연산을 수행할 수 있을 때, '같은 팀 여부 확인' 연산에 대한 연산 결과를 출력하는 프로그램을 작성하시오.
입력
7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1
출력
NO
NO
YES
✅ 해설
- 전형적인 서로소 집합 찾기 알고리즘 응용 문제
- 같은 팀 여부를 확인할 땐 최상위 부모 노드를 비교
- 최상위 부모 노드가 같다면 같은 팀
- 최상위 부모 노드가 다르다면 다른 팀
- 자세한 내용은 https://evolving-developer-new.tistory.com/58 참고
💻 코드
# 학생 수와 연산의 개수 입력받기
n, m = map(int, input().split())
# 부모 노드를 저장하는 배열
parent = []
# 자기 자신으로 부모노드 초기화
for i in range(n+1):
parent.append(i)
# 최상위 부모 노드를 찾는 함수
def find_top(n):
if parent[n] != n:
parent[n] = find_top(parent[n])
return parent[n]
# 팀 합치기 함수
def merge(n1, n2):
n1_parent = find_top(n1)
n2_parent = find_top(n2)
if n1_parent > n2_parent:
parent[n1_parent] = n2_parent
else:
parent[n2_parent] = n1_parent
# 결과를 저장할 배열
result = []
# 명령 입력받기
for _ in range(m):
order, n1, n2 = map(int, input().split())
# 같은 팀 여부 확인
if order == 1:
# 최상위 부모 노드 비교
if find_top(n1) == find_top(n2):
result.append("YES")
else:
result.append("NO")
# 팀 합치기
else:
merge(n1, n2)
for i in result:
print(i, end=' ')