백준 13518 - 트리와 쿼리 9 (C++)

문제

• 문제 링크

  
   BOJ 13518 - 트리와 쿼리 9 
  

• 문제 요약

N개의 정점으로 이루어진 트리가 주어진다.
Q개의 쿼리에 대한 답을 출력해보자.

• 제한

TL : $2$ sec, ML : $512$ MB




$2 ≤ N ≤ 100,000$

$1 ≤ Q ≤ 100,000$

$1 ≤ E_w ≤ 1,000,000$

알고리즘 분류

• 트리(trees)
• 최소 공통 조상(lowest common ancestor)
• 오일러 경로 테크닉(euler_tour_technique)
• 오프라인 쿼리(offline_queries)
• mo's(mo's algorithm)

풀이

예제를 예시로 들고, 트리를 그려보자.



그럼 대략 위와 같은 모양으로 그려볼 수 있는데, 이때 $1$부터 $ETT$를 돌려보며 기록 순서를 작성해보면 다음과같이 써볼 수 있다.

$(1)in[1] -> (2)in[2]-> (3)out[2] -> (4)in[3] -> (5)in[5] -> (6)out[5] -> (7)in[6] -> (8)out[6] -> (9)in[7] ->$
$(10)out[7] -> (11)out[3] -> (12)in[4] -> (13)in[8] -> (14)out[8] -> (15)out[4] -> (16)out[1]$

여기서 임의의 정점 쌍 $(p, q)$ 에 대해 $(p == LCA(p, q)) || (q == LCA(p, q))$ 를 만족 하거나 / 아니거나 둘 중 하나로 볼 수 있다.

---

만족 한다면 ?

$Q_1 : (1, 5)$ 라 하고 이를 위에 써둔 기록에 올려보자.

$(1)in[1] -> (2)in[2] -> (3)out[2] -> (4)in[3] -> (5)in[5]$ 가 되며 유일한 경로에 포함되지 않는 정점( $2$ )들은 $in[], out[]$ 에 의해 상쇄되고 결국 경로 상의 정점들만 남게 된다.

똑같이 $Q_2 : (4, 8)$ 에도 적용해보면
$(12)in[4] -> (13)in[8]$ 가 되며 경로 상의 정점들만 남게 된다.

결과적으로 $ETT$로 메겨지는 번호에 의해 구간을 짤라낼 수 있으며 $(p == LCA(p, q)) || (q == LCA(p, q))$ 를 만족하는 모든 쿼리에 대해 $[in[p], in[q]]$의 구간 쿼리로 환원 시켜 처리할 수 있게된다.

만족 하지 않는다면 ?

$Q_1 : (2, 5)$ 라 하고 이를 위에 써둔 기록에 올려보자.

$(3)out[2] -> (4)in[3] -> (5)in[5]$ 가 되며 경로 상에 $LCA(2, 5)$를 제외한 모든 정점들이 등장한다.

똑같이 $Q_2 : (7, 8)$ 에도 적용해보면
$(10)out[7] -> (11)out[3] -> (12)in[4] -> (13)in[8]$ 가 되며 $LCA(7, 8)$를 제외한 경로 상에 모든 정점들이 등장한다.

결과적으로 $ETT$로 메겨지는 번호에 의해 구간을 짤라낼 수 있으며 $(p == LCA(p, q)) || (q == LCA(p, q))$ 를 만족하지 않는 모든 쿼리에 대해 $[out[p], in[q]]$ 의 구간 쿼리로 환원시켜 처리할 수 있게된다.
그리고 추가로 여기에 포함되지 않았던 $LCA(p, q)$를 추가로 포함해주는 것이다.

---

이때, 위에서의 $u, v$는 항상 $in[u] < in[v]$ 를 만족 하도록 일반성을 유지하자.
그럼 이제 업데이트가 없는 구간 쿼리 문제가 되었으므로 $mo$'$s$로 빠르게 처리해줄 수 있다.

문제의 쿼리를 수행하기 위해 정점의 수를 카운트할 배열( $NC[]$ ), 가중치를 카운트할 배열( $WC[]$ ) 각각을 만들어주고 구간의 변화에 맞게 쿼리에 대한 답을 뽑아내주자.

오프라인 처리에서는 별로 어려울 게 없으므로 간단히 짚고 넘어가겠다.
현재 들고 있는 구간에 대하여,

• 어떤 정점이 구간에 포함될 때

  포함되는 순간 해당 정점이 유일하면서(= 카운트가 $1$) 가중치 또한 유일하다면 $t++$. 또는 해당 정점이 유일하지 않게 되면서(= 카운트가 $2$) 유일한 가중치가 사라졌다면 $t--$.

• 어떤 정점이 구간에서 빠질 때

  빠지는 순간 해당 정점이 유일하게 되면서(= 카운트가 $1$) 가중치 또한 유일하게 됐다면 $t++$. 또는 유일했던 해당 정점이 빠지면서(= 카운트가 $0$) 유일했던 가중치도 사라졌다면 $t--$.

전체 코드

#include<bits/stdc++.h>
#define N 100001
using namespace std;
 
struct _Q { int s, e, i, k; }Q[N];
vector <int> Gr[N];
int in[N], out[N], rec[N * 2], dp[N][18];
int D[N], W[N], NC[N], WC[N * 10], an[N];
int n, q, i, j, t;
 
void Init1()
{
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n;
    for (i = 1; i <= n; cin >> W[i++]);
    for (q = 1; q < n; q++)
        cin >> i >> j, Gr[j].push_back(i), Gr[i].push_back(j);
}
int o, ro;
void LCAInit1(int p, int q)
{
    D[p] = q, in[p] = ++o, rec[++ro] = p;
    for (int& i : Gr[p])
        if (!D[i])
            dp[i][0] = p, LCAInit1(i, q + 1);
    out[p] = ++o, rec[++ro] = p;
}
void LCAInit2()
{
    for (j = 1; j < 18; j++)
        for (i = 1; i <= n; i++)
            dp[i][j] = dp[dp[i][j - 1]][j - 1];
}
int GetLCA(int p, int q)
{
    if (D[p] < D[q]) swap(p, q);
    int d = D[p] - D[q], o{};
    for (o = 0; d; d >>= 1, o++)
        if (d & 1)
            p = dp[p][o];
    if (p ^ q)
    {
        for (o = 17; o >= 0; o--)
            if (dp[p][o] ^ dp[q][o])
                p = dp[p][o], q = dp[q][o];
        p = dp[p][0];
    }
    return p;
}
void Init2()
{
    for (cin >> q, o = 0; o < q; o++)
    {
        cin >> Q[o].s >> Q[o].e, Q[o].i = o;
        _Q& i(Q[o]);
        if (in[i.s] > in[i.e]) swap(i.s, i.e);
        j = GetLCA(i.s, i.e); i.e = in[i.e];
        i.s == j ? i.s = in[i.s] : (i.s = out[i.s], i.k = in[j]);
    }
    sort(Q, Q + q, [](_Q p, _Q q)
        {
            p.s /= 400, q.s /= 400;
            return p.s ^ q.s ? p.s < q.s : p.e < q.e;
        });
}
void MosQ(int p, int q)
{
    int x(W[p]); NC[p] += q;
    if (NC[p] == 1)
        WC[x]++, t += WC[x] == 1;
    if (NC[p] == (q > 0 ? 2 : 0))
        WC[x]--, t -= !WC[x];
}
void Solve(int s, int e)
{
    for (i = 0; i < q; i++)
    {
        while (s > Q[i].s) MosQ(rec[--s], 1);
        while (s < Q[i].s) MosQ(rec[s++], -1);
        while (e > Q[i].e) MosQ(rec[e--], -1);
        while (e < Q[i].e) MosQ(rec[++e], 1);
        if (Q[i].k) MosQ(rec[Q[i].k], 1);
        an[Q[i].i] = t;
        if (Q[i].k) MosQ(rec[Q[i].k], -1);
    }
    for (i = 0; i < q; cout << an[i++] << '\n');
}
int main()
{
    Init1();
    LCAInit1(1, 1); LCAInit2();
    Init2();
    Solve(0, 0);
}

comment

한 $10$연 맞왜틀을 경험 했지만 $LCA$와 $mo$'$s$가 엮인 재밌는 문제였다.
이런 문제는 어떻게 만드는 걸까..