Rerooting Technique on Tree

이 게시글에서는 트리에서 할 수 있는 rerooting technique라는 것을 다뤄본다.

Rerooting?

자연수 N개의 노드로 구성되고 루트가 정해진 트리가 주어져 있고 배열에 다음과 같은 정보가 담기길 원한다고 해보자.

 

sub[i] = i번째 노드가 루트인 서브 트리에 대해 해당 트리에 속한 노드의 개수

 

이 배열을 완성하기 위해서 어떻게 해야 할까? 먼저 생각나는 단순한 방법으로는 각 노드마다 dfs를 돌아 노드의 개수를 세는 $O(N^{2})$ 방법을 떠올릴 수 있다.

 

하지만 당연하게도 이런 글을 쓴다는 것은 $O(N^{2})$ 시간 복잡도로 해결할 수 없도록 N의 크기가 $10^{5}$ 이상인 경우에 대해 푸는 방법을 알아보는 것이므로 위 방법은 쓸 수 없다.

 

게다가 각 노드가 루트라고 할 때 전체 트리에 대해 어떤 연산을 한 결과까지 요구되면 이를 어떻게 해야 할지 까마득하다.

 

우리는 $O(N)$ DFS로 서브 트리를 구축하고 rerooting technique을 이용해 주어진 다른 연산을 $O(N)$ DFS로 처리하여 총 두 번의 DFS로 해결하는 방법을 알아볼 것이다.

 

이름에서 추측하면 루트를 재설정하는 과정을 거친다고 생각해볼 수 있다. 어떻게 재설정한다는 걸까?

서브 트리 별 노드 개수 구하기

일단 아무 트리나 한번 그려보자.

예쁜 이진트리다

여기서 1번 노드부터 7번 노드까지 각 서브 트리에 속한 노드 개수를 어떻게 DFS 한 번으로 구할 수 있을까?

 

사실 이건 동적 계획법을 생각해보면 해답을 빨리 찾을 수 있다.

각각 트리의 노드를 구하는 과정을 메모이제이션하면 DFS를 돌면서 노드의 개수를 누적할 수 있게 된다.

 

한번 DFS 코드를 짜 보자!

int sub[200001];
bool visited[200001];
vector<int> tree[200001];

fill(sub, sub + 1 + 200000, 1);

void dfsdp(int cur)
{
  visited[cur] = 1;
  for (int i: tree[cur])
  {
    if (visited[i]) continue;
    dfsdp(i);
    sub[cur] += sub[i];
  }
}

 

N이 20만 일 때 dfsdp 함수로 dfs를 돌면서 자식들의 sub값을 더해주는 것을 볼 수 있다.

 

Base case는 리프 노드에 도달했을 때인데 std::fill로 sub배열을 전부 1로 초기화해서 base case를 처리한다.

 

이렇게 하여 각 서브 트리의 노드 개수를 $O(N)$으로 구하게 된다.

 

이로써 우리는 서브 트리의 노드 갯수를 구축하게 되었다.

Rerooting Technique

그럼 서브 트리를 구했으니 이를 이용해 rerooting technique가 필요한 문제를 풀어보자.

 

https://atcoder.jp/contests/abc220/tasks/abc220_f

F - Distance Sums 2

 

https://cses.fi/problemset/task/1133

CSES - Tree Distances II

 

두 문제가 출력 공백을 제외하면 완전히 같으므로 편한 플랫폼으로 들어가서 문제를 보면 된다.

 

요약하면 각 노드별로 자신을 제외한 노드까지 가는 거리의 합을 출력하는 것이다.

 

이를 풀기 위해서는 다음 아이디어가 필요하다.

 

노드 x의 자식인 노드 y로 이동한다고 해보자.

이때 노드 y가 포함되지 않은 경로들의 길이는 1씩 커지고 y가 포함된 경로들의 길이는 1씩 작아진다.

 

그렇다. 여기서 reroot라는 게 뭔지 감을 잡았을 것이다. 루트가 변하면서 일어나는 변화를 잡아야 하는 것이다.

 

그러면 각 노드에 대해 정답을 담은 배열을 ans라고 하자. 위 문제에서 자식으로 이동하면서 생기는 변화는 다음과 같이 계산할 수 있다.

 

$ans(y) = ans(x) + \{ N - sub(y) \} - sub(y) = ans(x) + N - 2 sub(y)$

 

$N - sub(y)$는 y가 포함되지 않은 경로의 길이가 각각 1씩 커짐을 표현하고 $-sub(y)$는 y가 포함된 경로의 길이가 각각 1씩 작아짐을 뜻하게 된다.

 

우리는 루트부터 리프까지 DFS로 내려가면서 위 식을 적용하면 $O(N)$으로 답을 구할 수 있게 된다!

 

하지만 아직 문제가 남아있다. $ans(root)$를 모른다는 것이다. 초기 값을 알아야 식을 점진적으로 적용하던가 말던가 하는데 초기 값을 정하지 않았다.

 

이는 서브 트리를 구축하는 과정을 일부 수정하여 해결할 수 있다.

 

dfsdp 함수의 인자에 이동한 거리를 넘겨주도록 하여 자식으로 이동할 때마다 $ans(root)$에 이동한 거리를 더해주면 서브 트리를 구축하면서 $ans(root)$를 구할 수 있다. ans(1)을 구해보도록 하자.

 

void dfs(int cur, int mv)
{
  visited[cur] = true;
  ans[1] += mv;
  for (int i: tree[cur])
  {
    if (visited[i]) continue;
    dfs(i, mv + 1);
    sub[cur] += sub[i];
  }
}

 

dfs에서 이동할 때마다 ans[1]에 이동한 거리 mv를 더해주는 것을 볼 수 있다.

 

우리는 이렇게 $ans(root)$까지 알게 되었으니 남은 건 다른 dfs 함수를 만들어서 위 식을 적용시키면 되는 것이다!!

 

void ansquery(int cur)
{
  visited[cur] = true;
  for (int i: tree[cur])
  {
    if (visited[i]) continue;
    ans[i] = ans[cur] + n - 2 * sub[i];
    // ans[i] = ans[cur] + (n - sub[i]) - sub[i];
    ansquery(i);
  }
}

 

이렇게 우리는 단 두 번의 DFS로 모든 정답을 알게 되었다. 이렇게 구한 ans를 출력 조건에 맞게 출력해주면 된다.

 

구현이 어려우면 앳코더 문제 Dintance sums 2를 해결하는 코드를 참조하자.

 

https://atcoder.jp/contests/abc220/submissions/26196485

Submission #26196485 - AtCoder Beginner Contest 220

 

지금까지 우리는 서브 트리를 $O(N)$에 구축하여 rerooting technique에 써먹는 방법을 알아보았다.

 

이렇게 루트의 변화를 영리하게 이용해야 문제는 이것 말고도 더 있을 수 있으며 더 고급 주제에 대한 내용이 코드포스 블로그에 기록되어 있다.

 

하지만 아직까지 그것을 다룰 필요는 없다 생각하여 더 적지 않으나 필요가 생기면 2편으로 돌아오도록 하겠다.

 

마지막으로 Rerooting technique를 이용해서 푸는 문제를 하나 더 소개하는 것으로 마치겠다.

 

https://atcoder.jp/contests/abc160/tasks/abc160_f

F - Distributing Integers