아쉬움이 남지만 (BOJ 18862)

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

정점 $x$ 의 부모 정점의 번호를 $p(x)$ , 정점 $x$ 의 높이를 $h(x)$ 라고 하자. 이 때 정점 $c$ 의 조상 정점 $x$ 에서 정점 $c$ 까지 점프했을 때의 높이 $j(x, c)$ 는 아래와 같이 계산된다.

$j(x, c)=\sum\limits_{i=0}^{depth(c)-depth(x)}{\frac{h(p^{i}(c))}{2^{i+1}}}+\frac{h(x)}{2^{depth(c)-depth(x)+1}}$

$j(p(x), c)-j(x, c)=\frac{h(p(x))-h(x)}{2^{depth(c)-depth(x)+1}}$

dfs를 통해 깊이 있는 정점부터 답을 구하면서, $dfs(x)$ 에서 $H[c]=\lfloor j(x, c)\rfloor$ 가 유지되도록 하자.

void dfs(int vertex) {
    for (int child : adj[vertex]) {
        if (child == p(vertex)) continue;
        dfs(child, vertex);
    }
    add(vertex, 2 * h(vertex));
    answer[vertex] = count[vertex];
}

이 함수에서 $count[x]$ 는 $x$ 의 조상 정점 중에서 점프로 $x$ 에 도달할 수 있는 정점이 있을 경우 $0$ 이고, 없을 경우 $x$ 에서 점프로 도달할 수 있는 정점의 개수가 된다. add 함수는 현재 정점부터 시작하여 서브트리 전체의 $H, count$ 배열을 갱신할 것이다. 편의상 모든 정점의 높이가 짝수가 되도록 하기 위해 높이에 2를 곱해주었다.

void add(int vertex, int height) {
    if (height == 0) return;
    H[vertex] += height;
    R[vertex] += height;
    for (int child : adj[vertex]) {
        if (child == p(vertex)) continue;
        if (count[child] == 0) {
            add(child, R[vertex] / 2);
            count[vertex] += count[child];
            count[child] = 0;
        }
        else if (H[vertex] >= H[child]) {
            add(child, (H[vertex] - H[child]) / 2);
            count[vertex] += count[child];
            count[child] = 0;
        }
    }
    R[vertex] %= 2;
}

$R$ 은 같은 값이 두 번 더해지는 것을 막기 위한 배열으로, 항상 $R[x]=H[x]\mod 2$ 이다. add 함수를 호출하면 $j(x, c)=H[c].(Z[p^1(c)])(Z[p^2(c)])\dots(Z[x])_{(2)}$ 가 된다는 것을 계산을 통해 알 수 있으므로, 이 알고리즘은 올바른 답을 구한다.

이 알고리즘의 시간복잡도를 계산해보자. 각 정점 $c$ 에 대해서, $add(c, h)$ ( $h>0$ )가 호출되는 횟수가 $O(\log X)$ 임을 보이면 전체 시간복잡도가 $O(N\log X)$ 임을 보일 수 있다. ( $X=max_{1\le i\le N}(h(i))$ ) 일단 $dfs(x)$ 에서 $add(c)$ 가 호출되려면 $x$ 가 $c$ 의 조상이어야 한다. $depth(c)-depth(x)\le\log X+5$ 인 $x$ 는 $O(\log X)$ 개이므로 $add(c)$ 의 호출도 $O(\log X)$ 회이다. $depth(c)-depth(x)>\log X+5$ 인 경우에는 $j(x, c)$ 의 변화량 총 합이 최대 $1$ 이다. 따라서 $H[c]=\lfloor j(x, c)\rfloor$ 의 변화량도 최대 $1$ 이고, $add(c)$ 의 호출도 최대 $1$ 회이다. 따라서 $add(c)$ 는 총 $O(\log X)$ 회 호출되며, 총 시간복잡도는 $O(N\log X)$ 가 된다.