백준 16587 - Hierarchical Structure (C++)

문제

• 문제 링크

  
   BOJ 16587 - Hierarchical Structure 
  

• 문제 요약

$N$개의 정점으로 이루어진 트리가 주어진다.
$Q$개의 경로 쿼리에 대해 알맞는 답을 출력하자.

• 제한

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




$1 ≤ N, Q ≤ 100,000$

$1 ≤ A_i ≤ 10^9$

알고리즘 분류

• 자료 구조(data structures)
• 트리(trees)
• heavy-light 분할(heavy-light decomposition)
• 세그먼트 트리(segment tree)
• 머지 소트 트리(merge sort tree)
• 정렬(sorting)
• 누적 합(prefix sum)
• 이분 탐색(binary search)

풀이

쿼리의 내용을 정리하면 다음과 같이 생각할 수 있다.

두 정점을 잇는 경로 위 정점들의 가중치를 $x_i$라 할 때, $l ≤ x_i ≤ r$인 $Σx_i$ 의 값을 구하여라.

만약 경로 위 정점들의 가중치들을 모두 정렬된 상태로 갖고 있다면, 이분 탐색을 이용해 $l, r$의 인덱스를 빠르게 찾아낼 수 있다.
또, $[l, r]$ 는 연속적인데 이들의 연속합을 빠르게 구해야 한다. 즉, 누적 합을 이용할 수 있을 듯 하다.

---

구체적으로, 아래 전처리를 진행 한다는 것이다.

• 경로 상 쿼리를 선형적으로 처리할 수 있도록 hld 로 트리를 분할하자.
• 처음 주어지는 $A_i$들을 머지 소트 트리 ( $seg$ ) 위에 모두 올리자.
• $seg$의 모든 노드들을 정렬하며, 동시에 같은 노드 번호를 갖는 누적합 트리( $pseg$ ) 역시 작업해두자.

이제 $l ≤ x_i ≤ r$ 인 $x$의 인덱스 범위를 이분 탐색을 이용해 $O(logN)$만에 찾고, 그 사이 합을 $pseg[]$를 이용해 $O(1)$로 구하면 쿼리당 $O(log^3N)$의 복잡도로 문제를 해결할 수 있다.

전체 코드

#include<bits/stdc++.h>
#define r int n, int s, int e, int p, int q
#define ll long long
#define N 100001
using namespace std;
 
vector <ll> Gr[N], G[N], pseg[1 << 18];
vector <int> seg[1 << 18];
int D[N], S[N], P[N], C[N];
int top[N], in[N];
ll n, q, o, i, j;
 
void Init()
{
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n >> q;
    for (i = 1; i <= n; cin >> C[i++]);
    for (int o(1); o < n; o++)
        cin >> i >> j, Gr[i].push_back(j), Gr[j].push_back(i);
}
void HldDfs1(int p)
{
    S[p] = 1;
    for (int i : Gr[p])
        if (!S[i])
        {
            D[i] = D[p] + 1;
            P[i] = p;
            HldDfs1(i);
            S[p] += S[i];
            G[p].push_back(i);
            if (S[i] > S[G[p][0]]) swap(G[p][0], G[p].back());
        }
}
void HldDfs2(int p)
{
    in[p] = ++o;
    for (int i : G[p])
        top[i] = i == G[p][0] ? top[p] : i, HldDfs2(i);
}
void SegInit(r)
{
    if (s > p || e < p) return;
    seg[n].push_back(q);
    if (s ^ e)
        SegInit(n << 1, s, (s + e) >> 1, p, q), SegInit(n << 1 | 1, ((s + e) >> 1) + 1, e, p, q);
}
ll SegQuery(r)
{
    if (s > q || e < p) return 0;
    if (s >= p && e <= q)
    {
        ll x = lower_bound(seg[n].begin(), seg[n].end(), i) - seg[n].begin();
        ll y = upper_bound(seg[n].begin(), seg[n].end(), j) - seg[n].begin();
        return (y ? pseg[n][y - 1] : 0) - (x ? pseg[n][x - 1] : 0);
    }
    return SegQuery(n << 1, s, (s + e) >> 1, p, q) + SegQuery(n << 1 | 1, ((s + e) >> 1) + 1, e, p, q);
}
ll HldQuery(int x, int y, ll t = 0)
{
    while (top[x] ^ top[y])
    {
        if (D[top[x]] < D[top[y]]) swap(x, y);
        t += SegQuery(1, 1, n, in[top[x]], in[x]);
        x = P[top[x]];
    }
    if (D[x] > D[y]) swap(x, y);
    return t + SegQuery(1, 1, n, in[x], in[y]);
}
void Solve()
{
    for (i = 1; i <= n; i++) SegInit(1, 1, n, in[i], C[i]);
    for (i = 0; i < (1 << 18); i++)
    {
        sort(seg[i].begin(), seg[i].end());
        for (o = j = 0; j < (int)seg[i].size(); j++)
            o += seg[i][j], pseg[i].push_back(o);
    }
    for (int x, y; q--;) cin >> x >> y >> i >> j, cout << HldQuery(x, y) << '\n';
}
int main()
{
    Init();
    HldDfs1(1), HldDfs2(1);
    Solve();
}

comment

hld + 세그를 이용한 가장 기본적인 쿼리 처리 문제가 P1이다.
푼 사람이 별로 없어 기여가 적어서 그런지 몰라도, 충분히 D5로 뛸만하다고 생각한다.