백준 25973 - 어지러운 트리 (C++)

문제

• 문제 링크

  
   BOJ 25973 - 어지러운 트리 
  

• 문제 요약

$N$개의 정점으로 이루어진 트리와 $Q$개의 쿼리가 주어진다.
각 쿼리마다 그에 맞는 처리를 진행하자.

• 제한

TL : $1$ sec, ML : $1024$ MB




$1 ≤ N, Q ≤ 200,000$

$2$번 쿼리는 적어도 하나 이상 주어진다.

알고리즘 분류

• 트리(trees)
• 최소 공통 조상(lowest common ancestor)
• 다이나믹 프로그래밍(dp)
• 트리에서의 다이나믹 프로그래밍(dp_tree)

풀이

매 쿼리마다 $dfs$ 돌려볼 순 없는 노릇이다.

루트가 $1$인 최초 시점을 기준으로 모든 경우의 수를 전처리 후,
$x$와 현재 루트 정점을 가지고 뭐 어떻게 잘 요리해서 쿼리당 $O(logN)$ 으로 응답할 것을 생각하자.

경우의 수는 서브 트리의 사이즈 $(S[])$ 로써 계산할 수 있으니 $LCA$를 위한 전처리를 하며 임의의 정점에 대한 경우의 수 전처리도 동시에 진행할 수 있다.

즉, 다음의 추가적인 정보를 전처리 하자는 것이다.

$Sdp[i] :$ 루트가 $1$인 시점을 기준으로 $x == i$ 일 때 $(a, b)$의 쌍으로 가능한 경우의 수.

---

$1$번 쿼리는 간단하다. 말그대로 현재 루트를 $x$로 변경해 준다.

$2$번 쿼리는 $3$가지로 분류하여 생각해보자.

1. $root == x$ 일 때.
2. $LCA(root, x) == x$ 일 때.
3. $root$와 $x$가 아무런 연고가 없을 때.

---

$1$의 경우.

아래 내용에서의 $S[], Sdp[]$는 루트가 $1$일 때를 기준으로 전처리 한 것이다. 이를 항상 명심하여 식을 세워보자.

기본적으로 포함 되는 값은, $Sdp[x]$와 $S[x]$ $-$ $1$ 이다. $($후자는 $x$와 그 서브 트리 내 노드를 고르는 경우의 수$)$
그리고 $x$가 루트가 되면서 반전 된 경우의 수를 더해줘야 한다.

이는 위에서의 $S[x]$ $-$ $1$ 과 나머지 노드들인 $N$ $-$ $S[x]$ 를 곱한 값, $N$ $-$ $S[x]$ 를 더해주면 된다.

소거할 항을 소거하고 식을 정리해 보면

$r = Sdp[x] + S[x] - 1 + (S[x] - 1) * (N - S[x]) + N - S[x]$
$r = Sdp[x] + S[x] - 1 + NS[x] - S[x]^2 - N + S[x] + N - S[x]$
$r = Sdp[x] - 1 - S[x] * (S[x] - N - 1)$

$2$의 경우.

가장 까다로운 경우이다. 포함 및 배제할 값들을 정확히 추리기 위해 $o$를 $x$ 바로 밑까지 끌어올려 주고, 이 노드를 $y$라 하자.

기본적으로 포함 되는 값은, $Sdp[x]$와 $N - S[y] - 1$ 가 된다. $($후자는 $x$와 바뀐 서브 트리 내 노드를 고르는 경우의 수$)$
그리고 $x$의 $S[y]$를 제외한 기존 하위 노드들과 새로이 $x$의 하위 노드가 된 노드들간의 경우의 수를 더해줘야 한다.

$(S[x] - S[y] - 1) * (N - S[x])$

처음 포함한 $Sdp[x]$는 $S[y]$의 노드들이 고려된 값이다.
관계가 뒤집힘에 따라 발생하는 차이만큼 빼주어야 하므로 $S[y] * (S[x] - S[y] - 1)$ 를 빼주면 된다.

소거할 항을 소거하고 식을 정리해 보면

$r = Sdp[x] + N - S[y] - 1 + (S[x] - S[y] - 1) * (N - S[x]) - S[y] * (S[x] - S[y] - 1)$
$r = Sdp[x] + N - S[y] - 1 + NS[x] - S[x]^2 - NS[y] + S[y]S[x] - N + S[x] - S[y]S[x] + S[y]^2 + S[y]$
$r = Sdp[x] - 1 + S[y] * (S[y] - N) - S[x] * (S[x] - N - 1)$

$3$의 경우.

간단하다.
루트가 어떤 노드이건 상관 없이 $x$와 $x$의 서브 트리로만 구성된다.

$r = Sdp[x] + S[x] - 1$

---

좀 주저리주저리 한 것 같은데.. 적당한 사이즈의 트리를 만들어 놓고,
이런저런 예제를 만들어 보며 직접 식을 유도해 보는 것을 추천한다.

그 과정에서 진정 문제를 제대로 바라보게 될 테니 말이다.

전체 코드

#include<bits/stdc++.h>
#define N 200001
using namespace std;
 
vector <int> Gr[N];
int n, q, D[N], dp[N][19];
long long r, S[N], Sdp[N];
 
void Init()
{
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n >> q;
    for (int i, j, o{}; ++o < n;)
        cin >> i >> j, Gr[i].push_back(j), Gr[j].push_back(i);
}
void LCAInit(int p, int q)
{
    D[p] = q; S[p] = 1;
    for (int i : Gr[p])
        if (!D[i])
        {
            LCAInit(i, q + 1);
            Sdp[p] += S[i] * (S[p] - 1);
            dp[i][0] = p, S[p] += S[i];
        }
}
void SparseTableDP()
{
    for (int j(1); j < 19; j++)
        for (int i(1); i <= n; i++)
            dp[i][j] = dp[dp[i][j - 1]][j - 1];
}
int GetLCA(int p, int q, int k)
{
    if (D[p] < D[q]) swap(p, q);
    int d(D[p] - D[q] - k);
    for (int i{}; d; d >>= 1, i++)
        if (d & 1)
            p = dp[p][i];
    if (!k && p ^ q)
    {
        for (int i(18); i >= 0; i--)
            if (dp[p][i] ^ dp[q][i])
                p = dp[p][i], q = dp[q][i];
        p = dp[p][0];
    }
    return p;
}
void Solve(int rt = 1)
{
    for (int i, x, y; q--;)
        if (cin >> i >> x, i & 1) rt = x;
        else
        {
            if (rt == x)
                r = -S[x] * (S[x] - n - 1);
            else if (x == GetLCA(x, rt, 0))
            {
                y = GetLCA(x, rt, 1);
                r = S[y] * (S[y] - n) - S[x] * (S[x] - n - 1);
            }
            else
                r = S[x];
            cout << r + Sdp[x] - 1 << '\n';
        }
}
int main()
{
    Init();
    LCAInit(1, 1);
    SparseTableDP();
    Solve();
}

comment

$LCA$와 좀 더 친해질 수 있는 좋은 문제라고 생각한다.