트리 T는 아래 그림 1과 같은 구조를 가지고 있으며 원은 ‘정점’이라 하고, 정점과 정점을 연결하는 선을 ‘에지’라 한다. 특히 가장 위에 위치한 정점을 ‘루트’라 하는데 오직 하나만 있다. N개의 정점들은 숫자 1부터 N으로 표현하고 루트는 항상 1이다.
두 정점 v와 w를 연결하는 경로는 정점들의 순서리스트 (v0, v1, ..., vm)로, 정점 vi와 vi+1은 에지로 연결되고 v0 = v, vm = w이다. 트리에서는 임의의 두 정점 v와 w 사이에 항상 두 정점을 연결하는 경로가 오직 하나만 존재한다. 예를 들어, 그림 1에서 정점 3과 11 사이의 유일한 경로는 (3, 4, 1, 7, 11)이다.
각 정점 v에서 루트 r과 연결하는 유일한 경로 P에 대해서 정점 v와 에지로 연결된 정점 중에서 P상에 있는 정점을 v의 ‘부모 정점’이라고 한다. 예를 들어, 그림 1에서 4, 7, 9의 부모 정점은 1이고, 2와 11의 부모 정점은 7이다.
트리 T에서 어떤 두 정점을 연결하는 에지를 제거하면 그 두 정점 외에도 경로가 존재하지 않는 정점 쌍이 있을 수 있다. 여러분은 “정점 v와 w를 연결하는 경로가 존재하는가?”와 같은 질의에 답해야 한다. 예를 들어, 그림 1에서 7과 11 사이의 에지를 제거하면 8과 5를 연결하는 경로는 존재하지 않는다.
트리 정보가 주어지고, 에지의 제거 정보와 질의가 임의의 순서로 주어질 때, 작업을 순서대로 수행하며 질의에 대한 답을 출력하는 프로그램을 작성하시오.
표준 입력으로 다음 정보가 주어진다. 첫 번째 줄에는 트리의 정점의 개수와 질의의 개수를 나타내는 두 정수 N과 Q (1 ≤ N, Q ≤ 200,000)가 주어진다. 다음 N-1개의 줄의 i번째 줄에는 정점 i+1의 부모 정점을 나타내는 정수 a가 주어진다 (1 ≤ a ≤ N). 다음 (N-1)+Q개의 줄 중에서 N-1개는 (1)의 형태로, Q개는 (2)의 형태로 주어진다. (1) 두 정수 x와 b가 주어진다(x = 0, 2 ≤ b ≤ N). 이것은 b의 부모 정점과 b를 연결하는 에지를 제거함을 의미한다. 각 줄의 b는 모두 다르다. (2) 세 정수 x, c, d가 주어진다 (x = 1, 1 ≤ c, d ≤ N). 이것은 c와 d를 연결하는 경로가 존재하는 지 묻는 질의를 의미한다.
출력
표준 출력으로 질의에 대한 답을 순서대로 Q개의 줄에 출력한다. 각 줄마다 경로가 존재하면 YES를 아니면 NO를 출력한다.
입력 예시
3 3
1
1
1 2 3
0 3
1 2 3
1 1 2
0 2
출력 예시
YES
NO
YES
풀이
임의의 두 점의 연결 여부를 구하려면 유니온 파인드를 사용하면 된다. 경로가 단축된 parent[node]는 그 노드가 속한 트리의 루트를 가리키게 되므로, 두 노드의 루트 노드가 같다면 YES, 그렇지 않다면 다른 트리에 속했으므로 NO를 출력하는 것이 첫 번째 포인트가 되겠다.
그런데, 여기서 주어지는 쿼리는 경로를 끊는 작업을 수행한다. 즉 유니온의 반대 케이스라고 볼 수 있다. 이 경우를 순차적으로 처리하려면 여러 가지 제약이 따른다(그대로 수행하자니 트리의 하위 노드들의 연결성이 보장되지 않고 - 즉 아직도 끊어지기 전 루트 노드를 가리킴 - 경로가 끊어진 지점부터 트리를 업데이트하더라도 메모리 및 시간적 제약에 걸린다) 즉 우리는 발상의 전환을 할 필요가 있다.
모든 쿼리를 수행하면 결과적으로 N-1개의 단일 노드가 되고, 이를 거꾸로 수행한다면 경로를 끊는 작업은 두 지점의 경로를 잇는 작업으로 변환된다. 즉 쿼리를 거꾸로 수행하여, 유니온으로 경로를 잇고 Find로 두 노드의 연결성을 검사하면 유니온 파인드만으로도 모든 경우를 올바르게 탐색할 수 있다. 이렇게 쿼리를 순차적으로 처리하는 것이 아니라 필요에 따라 순서를 바꿔가며 처리하는 기법을 오프라인 쿼리라고 한다. 쿼리를 거꾸로 수행했으므로 정답 출력을 다시 역순으로 출력해야 한다는 점에 주의하자.
풀이 코드
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline
N, Q = map(int, input().split())
union_target = [1]*(N+1)
parents = list(range(N+1))
for i in range(2, N+1) :
ui = int(input())
union_target[i] = ui
ans = list()
def find(a) :
if a == parents[a] :
return a
parents[a] = find(parents[a])
return parents[a]
def union(a, b) :
parents[b] = a
queries = [list(map(int, input().split())) for _ in range(N-1+Q)]
for q in reversed(queries) :
if q[0] == 0 :
b = q[1]
a = find(union_target[b])
union(a, b)
else :
a, b = q[1:]
ans.append("YES" if find(a) == find(b) else "NO")
print(*reversed(ans), sep = "\n")
