[이코테] 팀 결성 - 이해까지 도와주는 풀이

문제

학교에서 학생들에게 0번부터 N번까지의 번호를 부여했다. 처음에는 모든 학생이 서로 다른 팀으로 구분되어, 

총 N + 1 개의 팀이 존재한다. 이때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.

 

1. '팀 합치기' 연산은 두 팀을 합치는 연산이다.

2. '같은 팀 여부 확인' 연산은 특정한 두 학생이 같은 팀에 속하는 지 확인하는 연산이다.

 

선생님이 M개의 연산을 수행할 수 있을 때, '같은 팀 여부 확인' 연산에 대한 연산 결과를 출력하는 프로그램을 작성하시오.

 

입력 조건

 

1. 첫째 줄에 N, M이 주어진다. M은 입력으로 주어지는 연산의 개수이다. (1 <= N, M <= 100,000)

 

2. 다음 M개의 줄에는 각각의 연산이 주어진다.

 

3. '팀 합치기' 연산은 0 a b 형태로 주어진다. 이는 a번 학생이 속한 팀과 b번 학생이 속한 팀을 합친다는 의미이다.

 

4. '같은 팀 여부 확인' 연산은 1 a b 형태로 주어진다. 이는 a번 학생과 b번 학생이 같은 팀에 속해 있는지를 확인하는 연산이다.

 

5. a와 b는 N 이하의 양의 정수이다.

 

출력조건

 

'같은 팀 여부 확인' 연산에 대하여 한 줄에 하나씩 YES 혹은 NO로 결과를 출력한다.

<입력 예시>
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

 

 

 

풀이

 

이 문제는 한국어로 서로소라고 하며 Disjoint Set(또는 Union-Find) 자료 구조를 사용하여 풀 수 있는 전형적인 문제 중 하나입니다. 

 

여러 개의 집합(set)이 주어지고, 이 집합들 간의 공통 원소가 없는지 판단하는 문제입니다.

즉, 각 집합은 서로 중복된 원소를 가지지 않아야 합니다.

 

N과 M의 범위가 모두 최대 100,000입니다. 따라서 경로 압축 방식의 서로소 집합 자료구조를 이용하여 시간 복잡도를 개선해야 합니다.

다양한 알고리즘 - 서로소 집합(경로 압축 기법)

 

Disjoint Set은 여러 개의 집합을 관리하고, 각 집합의 원소들이 서로 연결되어 있는지 여부를 효율적으로 확인하는 데 사용됩니다.

문제에서 설명한 내용을 간단히 요약하면 다음과 같습니다.

 

1. 학생들은 처음에 각자 고유한 팀을 가지고 있으며, 팀 번호는 0부터 N까지입니다.
2. 선생님은 두 가지 종류의 연산을 수행할 수 있습니다.
   - '팀 합치기' 연산: 두 팀을 합쳐 하나의 팀으로 만듭니다.
   - '같은 팀 여부 확인' 연산: 두 학생이 같은 팀에 속하는지 확인합니다.
3. '같은 팀 여부 확인' 연산이 수행될 때마다 결과를 출력해야 합니다.

이 문제를 해결하는 일반적인 접근 방법은 다음과 같습니다

1. 초기에 각 학생을 자기 자신을 가리키는 별개의 집합으로 초기화합니다. 즉, 학생 번호와 팀 번호가 일치하는 집합을 생성합니다.

2. '팀 합치기' 연산을 수행할 때, 두 학생이 속한 두 집합을 합칩니다. 이를 위해 Union 연산을 사용합니다.

3. '같은 팀 여부 확인' 연산을 수행할 때, 두 학생이 같은 팀에 속하는지 확인하기 위해 Find 연산을 사용합니다. Find 연산은 특정 원소가 속한 집합을 찾아주는 연산입니다.

4. '같은 팀 여부 확인' 연산의 결과를 출력합니다.

 

소스코드

# 루트 노드를 찾는 함수
def find_parent(parent, x):
    # 루트 노드가 아니라면, 루트 노드를 찾을 때까지 재귀적으로 호출
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])  # 경로 압축(Path Compression)을 통해 루트 노드를 갱신
    return parent[x]

# 두 집합을 합치는 함수
def union_parent(parent, a, b):
    a = find_parent(parent, a)  # a의 루트 노드를 찾음
    b = find_parent(parent, b)  # b의 루트 노드를 찾음

    if a < b:
        parent[b] = a  # a가 더 작은 값을 가진 경우, b의 루트를 a로 설정
    else:
        parent[a] = b  # b가 더 작거나 같은 값을 가진 경우, a의 루트를 b로 설정

n, m = map(int, input().split())
parent = [0] * (n + 1)

for i in range(n + 1):
    parent[i] = i  # 초기에는 각 원소가 자신을 가리키도록 초기화

for i in range(m):
    oper, a, b = map(int, input().split())
    if oper == 0:
        union_parent(parent, a, b)  # '팀 합치기' 연산을 수행하여 a와 b가 속한 팀을 합침
    elif oper == 1:
        if find_parent(parent, a) == find_parent(parent, b):
            print("YES")  # '같은 팀 여부 확인' 연산을 수행하여 a와 b가 같은 팀이면 "YES" 출력
        else:
            print("NO")  # a와 b가 다른 팀이면 "NO" 출력

 

결과

7 8
0 1 3
1 1 7
NO
0 7 6
1 7 1
NO
0 3 7
0 4 2
0 1 1
1 1 1
YES

 

이 코드는 Disjoint Set을 사용하여 각 학생이 어떤 팀에 속해 있는지 유지하고, '같은 팀 여부 확인' 연산에 대한 결과를 출력하는 예시입니다. Disjoint Set은 연산 시간이 상수 시간에 가깝기 때문에 매우 효율적으로 작동합니다.