My code
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(1, n + 1):
parent[i] = i
for i in range(m):
a, b, c = map(int, input().split())
if a == 0:
union_parent(parent, b, c)
else:
if find_parent(parent, b) == find_parent(parent, c):
print('YES')
else:
print('NO')
Answer
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())
if oper == 0:
union_parent(parent, a, b)
elif oper == 1:
if find_parent(parent, a) == find_parent(parent, b):
print('YES')
else:
print('NO')
Code Review
이번 문제는 팀 합치기, 같은 팀 여부 확인 연산을 M 번 수행할 때 같은 팀 여부 확인 연산에 대한 연산 결과를 출력하는 문제였습니다. 팀 합치기는 union 연산, 같은 팀 여부 확인은 find 연산을 이용하여 구할 수 있습니다. 먼저 같은 팀 여부 확인을 위한 find_parent 함수를 만들어주어 각 노드의 루트 노드를 찾을 수 있도록 해주었습니다. 만약 해당 노드의 parent 배열의 값이 노드의 값과 같지 않다면 루트 노드가 아니라는 뜻이므로 다시 한번 find_parent 함수를 재귀적으로 호출해주었고 이때 parent 배열에 들어있던 값을 노드 값으로 넣어주었습니다. 그리고 만약 루트 노드를 찾았다면 그 배열의 값을 return 하도록 해주었습니다.
다음으로 팀 합치기를 위한 union_parent 함수를 만들어주어 a의 루트 노드, b의 루트 노드를 찾아준 뒤 만약 a의 루트 노드가 b 보다 작다면 b 의 루트노드를 a의 값으로 갱신시켜주도록 하여 두 노드를 같은 루트 노드를 같도록 합쳐주었습니다.
이처럼 두 가지 연산을 위한 함수를 만든 뒤 먼저 초기값으로 루트 노드를 자기 자신으로 갖도록 parent 배열을 초기화 시켜줍니다. 그런 다음 세 값을 입력 받으며 만약 첫 번째 값이 0이라면 합치기 연산을 수행하도록, 1 이라면 같은 팀 여부를 확인하도록 해주었습니다. 즉, 0이라면 union_parent 함수를 실행시켜주어 두 노드를 같은 집합으로 합쳐주었습니다. 그리고 1이라면 두번째, 세번째 값에 대해 find_parent 함수를 실행시켰을 때 두 값이 같다는 것은 루트 노드가 서로 같다는 뜻이며 이는 같은 집합에 속한다는 것을 의미합니다. 따라서 같은 팀이라는 것을 뜻하므로 YES 를 출력하도록 해주었고 그렇지 않은 경우에는 NO 를 출력하도록 해주었습니다.
댓글