백준 5916번 - 농장 관리

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

5916번: 농장 관리

HLD 풀이가 정해이지만 LCA + 세그트리로 해결할 수 있습니다.

 

필요한 배경지식

• 최소 공통 조상(LCA)
• 세그먼트 트리
• 오일러 투어 테크닉

 

문제 해결 방법

우선 N개의 노드와 N-1개의 간선이 존재하므로 농장은 트리 형태의 자료구조를 가진다는것을 확인할 수 있습니다.

 

이 문제에서 조심해야 할 부분으로는 P와 Q쿼리 모두 노드가 아닌 간선의 값을 처리해야 하기 때문에 모든 간선이 자식노드에 속해있다고 생각하고 문제를 풀어야 합니다.

또한 Q쿼리의 내용을 보면, 출력하는 값으로 경로가 아닌 u번 농장과 v번 농장 사이의 '도로'라고 나와있으므로 포인트 쿼리를 출력해야 합니다.

 

이 문제에서처럼 구간 업데이트(P), 포인트 쿼리(Q)를 필요로 할때 구간 업데이트를 단순 연산(O(1))으로 처리할 수 있는 방법이 존재합니다. (https://algoshipda.blogspot.com/2017/05/lazy-propagation.html)

 

트리에서도 위의 내용을 적용시킬 수 있습니다.

n번노드의 간선크기가 k이고 k를 n노드부터 루트노드까지의 간선마다 더해질 값이라고 한다면,

다음과 같은 방법으로 P쿼리와 Q쿼리를 처리할 수 있습니다.

 

1. P u v : u농장과 v농장 도로에는 +1을 업데이트, LCA농장의 도로에는 -2를 업데이트
2. Q u v : u와 v사이의 도로를 포함하는 농장이 u라고 하면, 농장u와 u의 자식노드에 해당하는 모든 농장의 도로에 존재하는 나무들의 총합

Q쿼리에서, 트리에서의 구간합을 구해야 하므로 오일러 경로 테크닉과 세그먼트 트리를 필요로 합니다.

 

 

 

전체 코드

#include<bits/stdc++.h>
#define MAX 100000
using namespace std;

int N, M, maxPower = (int)floor(log2(MAX));
int ac[MAX+1][21]{};
int dpt[MAX+1]{};
vector<int> adj[MAX+1]{};
int st[MAX+1]{}, ed[MAX+1]{}; // 오일러 투어 in, out
int tree[MAX<<2]{};

int num = 0;
void dfs(int now, int parent) {
    st[now] = ++num;
    dpt[now] = dpt[parent] + 1;
    ac[now][0] = parent;

    for (int i = 1; i <= maxPower; i++)
        ac[now][i] = ac[ac[now][i-1]][i-1];

    for (auto next : adj[now]) {
        if (next == parent) continue;
        dfs(next, now);
    }
    ed[now] = num;
}

int lca(int x, int y) {
    if (dpt[x] > dpt[y]) swap(x, y);
    for (int i = maxPower; i >= 0; i--) { 
        if (dpt[y] - dpt[x] >= (1 << i)) 
            y = ac[y][i];
    }
    if (x == y)return x;
    for (int i = maxPower; i >= 0; i--) {
        if (ac[x][i] != ac[y][i]) {
            x = ac[x][i];
            y = ac[y][i];
        }
    }
    return ac[x][0];
}

void modify(int n, int s, int e, int idx, int diff){
    if(idx<s || idx>e) return;
    tree[n] += diff;
    if(s==e) return;
    int mid = (s+e)>>1;
    modify(n<<1, s, mid, idx, diff);
    modify(n<<1|1, mid+1, e, idx, diff);
}

int query(int n, int s, int e, int l, int r){
    if(l>e || r<s) return 0;
    if(l<=s && e<=r) return tree[n];
    int mid = (s+e)>>1;
    return query(n<<1, s, mid, l, r) + query(n<<1|1, mid+1, e, l, r);
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin>>N>>M;
    int a,b;
    for(int i=2; i<=N; i++){
        cin>>a>>b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    dfs(1, 0);

    char ch;
    while(M--){
        cin>>ch;
        cin>>a>>b;
        int next;
        if(ch=='P') {
            next = lca(a,b);
            modify(1,1,N, st[next],-2);
            modify(1,1,N, st[a],1);
            modify(1,1,N, st[b],1);
        } else {
            if(st[a] < st[b]) next = b;
            else next = a;
            cout<<query(1,1,N, st[next], ed[next])<<'\n';
        }
    }
}