BOJ 18189 참 어려운 문제

https://www.acmicpc.net/problem/18189

문제 설명

i번 정점의 색이 A[i]인 크기가 N인 트리가 주어질 때,

색깔이 같은 두 정점을 고르면, 이 둘은 항상 조상-자식 관계가 아니다

의 조건을 만족하는 루트가 될 수 있는 정점을 모두 찾아라.

풀이

r을 루트로 하는 트리에서 모든 정점 v에 대해 색이 A[v]인 정점이 v의 서브트리에 하나만 있다 면 r은 루트가 될 수 있다.

이 조건은 서브트리를 연속한 구간으로 바꿔서 빠르게 확인할 수 있다.

v에 들어오는 시간을 in[v], 나가는 시간을 out[v]라고 했을 때,
시간이 [in[v], out[v]]에 있는 정점들은 v의 서브트리에 속해있다.

색이 c인 정점을 방문하는 시간들을 time[c]에 저장하면 특정 시간 구간에 색이 c인 정점이 몇 개나 있었는지 알아낼 수 있다.

서브트리를 이용하기 위해 임의의 정점을 루트로 삼는다.

빨간 트리는 r의 서브트리이고, 파란 트리는 나머지 노드와 r을 가지고 구성한 트리이다.

r을 루트로 하는 트리를 생각해보면
r을 제외하고 빨간 트리와 파란 트리에서 각각 정점을 뽑았을 때, 조상-자식 관계가 아니므로
빨간 트리와 파란 트리에서 r이 루트가 될 수 있다면 전체 트리에서도 r은 루트가 될 수 있다.

어떤 정점 r이
빨간 트리에서 루트가 될 수 있는지 여부를 d[r],
파란 트리에서 루트가 될 수 있는지 여부를 p[r]라고 하자.

다음 조건을 만족한다면 d[r]은 true이다.

r의 자식들 x에 대하여, d[x]가 모두 true이다
r의 서브트리에 A[r]과 같은 색인 정점이 r밖에 없다

다음 조건을 만족한다면 p[r]은 true이다.

p의 자식이 r, c1, ..., ci일 때,
p[p]는 true이다
d[c1] = ... = d[ci] = true이다
r의 서브트리 바깥에 색이 A[r]과 같은 색인 정점이 없다.
c1, ... , ci의 서브트리에 A[p]와 같은 색인 정점이 없다

d[i] = p[i] = true인 i는 루트가 될 수 있다.