[백준 13896번] Sky Tax

1. LCA 알고리즘에서 2^0 번째 조상을 sparse table 에 저장하고,

각 정점의 depth 를 저장하기 위해 DFS 탐색을 수행하는데,

이 때, 자신을 포함한 모든 하위 노드의 갯수를 저장해둔다.

 

2. 1번 쿼리의 답은 아래와 같이 3가지 케이스로 나눌 수 있다.

1. R 과 U 가 같은 정점인 경우 = 답 : N(전체 정점 수)

2. LCA(R, U) 와 U 가 같은 정점인 경우 = 답 : N(전체 정점 수) - U의 자식들 중에서 R을 자식으로 둔 노드가 저장해둔 모든 하위 노드 갯수

3. LCA(R, U) 가 U 가 아닌 경우(U 보다 depth가 낮다) = 답 : U를 포함한 U의 모든 하위 노드의 갯수

 

핵심 부분이 되는 코드는 아래 추가해둔다.

 

LCA 의 DFS 탐색을 돌면서 cnt 배열에는 자신을 포함한 전체 하위 노드의 갯수를 저장

private static int dfs(int now, int p) {
    int num = 1;
    for (int next : adjList[now]) {
        if (next == p) continue;
        parent[0][next] = now;
        depth[next] = depth[now] + 1;
        num += dfs(next, now);
    }
    return cnt[now] = num;
}

 

쿼리 수행은 아래와 같이

LCA(R, U) == U 인 경우

R 에서 R 과 U 의 depth 차이 - 1 만큼 올라가면, U 의 바로 아래 자식이 된다.

그 자식이 가진 "자신을 포함한 전체 하위 노드의 갯수" 를 N (전체 정점 수) 에서 빼주면 된다.

if (s == 0) { // move capital
    R = u;
}
else { // compute number of cities
    bw.write("\n");
    int lca = lca(u, R);
    int ans;
    if (R == u) {
        ans = N;
    }
    else if (lca == u) {
        int a = move(R, depth[R] - depth[u] - 1);
        ans = N - cnt[a];
    }
    else {
        ans = cnt[u];
    }
    bw.write(ans + "");
}