BOJ 18803 방역

정점이 N개인 트리가 주어질 때, 정점들을 적당히 지워 어떤 임의의 두 점 사이의 거리도 K를 넘지 않도록 하게 지우는 경우의 수를 세는 문제이다. N, K≤300000이다.

 

#1. O(NK^2)

Tree DP를 통해 세제곱 시간에 문제를 풀어낼 수가 있다. DP[u][a]를 u의 서브트리에 대해, 움직일 수 있는 정점의 최소 깊이가 a이도록 정점을 제거할 수 있는 경우의 수라 정의하자. 당연히 어떤 두 점 사이의 거리도 K를 넘으면 안된다는 조건을 만족하도록 제거하는 것이다. 편의상 모든 경로가 K이하인 것으로 조건을 바꾸고, 입력에서 주어진 K에 -1을 해주도록 하겠다.

 

이러한 조건 하에서 상태전이는 꽤 직관적이다. 어떤 임의의 정점에 대해 u의 자식 노드를 v1, v2라 하면,

 

DP[u][max(a, b)+1]=∑(DP[v1][a]*DP[v2][b]) (단, a+b≤K)

 

정의상 정점u가 지워질 때의 subtree u에 대한 DP는 DP[u][-1]이 되는데... 일단 납득하고 넘어가자. 어차피 정해의 DP는 여기까지 고려할 것이다.

 

위 상태전이에서는 DP[u][-1]의 값을 갱신하지 못하므로, 이 경우에 대해서만 따로 전이를 처리해주자.

 

DP[u][-1]=(DP[v1][-1]+...+DP[v1][K-1])*(DP[v2][-1]+...+DP[v2][K-1])

 

참고로, 리프 정점들에 대해서는 DP[u][-1]=DP[u][0]=1으로 초항을 설정해주면 된다.

 

한번 합쳐주는데 O(K^2)이고, 합쳐주는 횟수는 O(N)이므로 O(NK^2) 풀이가 완성된다.

 

#2. O(NKlogK)

위의 DP 전이식을 잘 살펴보면, 적절히 구간 쿼리와 구간 업데이트로 바꿔서 같은 식을 표현할 수 있다는 것을 알 수 있다. 식을 잘 정리해주면 결국 우리가 필요한 것은 구간의 합과, 구간에 어떤 수를 곱해주는 쿼리임을 알 수 있으므로, 각 정점마다 Segment tree with Lazy Propogation을 들고 있으면 한번 합쳐주는데 O(KlogK)로 충분함을 알 수 있다.

 

#3. O(NlogNlogK)

한번 합쳐주는데 드는 시간복잡도를 조금 더 구체적으로 생각해보자. 정점 x의 높이를 H[]라 하자. 이때 subtree u와 subtree v를 합치는데 드는 시간복잡도는 O(min(H[u], k)log(min(H[v], k)))이다. 이는 양쪽 높이 중 작은 쪽의 높이에 의존적이므로, 높이에 대한 Small To Large 을 해주면 O(NlogNlogK)에 풀 수 있다.

 

이때 메모리를 위해 각 노드마다 포인터를 활용한 동적 Segment Tree를 갖고 있어야 함에 주의하자.

 

또한, 가장 크기가 큰 Segment Tree에 대한 전이를 할 때, DP[x][i]을 DP[x][i+1]로 한칸 미는 연산이 필요함을 알 수 있다. 이를 해결하기 위해 DP의 정의를 약간만 바꿔주자. DP[u][a]를 u의 subtree내에서 도달 가능한 최대 깊이가 a일 경우의 수이다. 즉, 기존의 DP[u][a]를 DP[u][d[u]+a]로 바꾼다고 생각하면 편하다. (d[u]는 정점 u의 깊이)

 

N, K가 30만이고, 포인터연산과 함께 Lazy Propogation을 해서 상수가 꽤 클 것 같지만, 놀랍게도 2초의 시간 제한 내에 실행된다.

 

#4. O(NlogK)

놀랍게도 위의 알고리즘의 시간복잡도는 O(NlogNlogK)이 아니라 O(NlogK)임을 증명할 수 있다. 서브트리의 크기에 대한 Small to Large가 아니라 높이에 대한 Small To Large여서 logN이 떼지는 것이다. 이에 대한 증명은 Cubelover님의 블로그에서 잘 찾아볼 수 있다.