BOJ 8145 - Megalopolis 링크
(2023.04.17 기준 P2)

문제

n개의 마을이 n-1개의 비포장 도로로 연결되어 있다.
다음과 같은 쿼리를 처리.
1. A a b : a와 b를 잇는 도로를 포장한다.
2. W a : 1번 마을에서 a번 마을로의 경로에서 비포장 도로의 개수 출력

알고리즘

마을은 트리 형태로 연결되어 있고, 간선의 정보가 바뀌며 각 구간의 정보를 업데이트 해야 한다. n이 최대 250,000이므로 lca를 이용하여 하나씩 구하면 TLE가 난다. 그러므로 HLD와 세그먼트 트리를 이용하자.

풀이

HLD는 처음 블로그에서 다룬다.
간단히 설명하자면 트리를 제일 무거운 간선, 나머지 가벼운 간선 집합으로 나누면서 체인을 만들어 가는 것이다. 그 체인마다 세그먼트 트리를 이용하여 정보를 즉각 업데이트한다고 생각하면 된다.
HLD에 대한 자세한 설명은 다음 블로그에서 참고하자.
1. 설명을 직관적으로 이해하기 정말 쉽게 해놓은 블로그
2. 좀 더 자세히 설명을 해주는 블로그

문제의 예제를 그림으로 나타내어 보자. 1을 루트로 하면
위와 같이 된다. 간선의 가중치는 비포장 도로임을 나타낸다.
정점도 아니고 간선의 가중치라 어려울 수도 있지만, 정점은 n, 간선은 n-1개다. 그러면 dfs ordering을 하였을 때, 부모와 자식 관계의 두 정점이 있을 건데, 자식의 정점에 간선의 정보를 입력하면 된다. 루트를 제외한 모든 정점은 부모가 하나씩 존재하기 때문이다.

이제 HLD를 하면위와 같이 된다. 빨간 체인이 무거운 간선들의 체인이다.
그럼 이제 각 체인별로 세그를 접목시키면 된다.
이 문제에선 도로 하나를 포장 및 비포장 도로의 개수 출력이므로 점 업데이트 및 구간합 쿼리를 구현하면 된다. 이는 EASY하므로 넘어가겠다.
아 그런데 이 문제에선 펜윅을 쓰지 않으면 TLE다. 주의하자.

코드 (2023.07.03 수정)

• C++

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

const int MAXN = 250000;

int n;

// 펜윅트리 (1-based index)
struct FW{
    int tree[MAXN + 1];

    void init();

    void update(int nd, int val){ // 점 업데이트
        while (nd <= n){
            tree[nd] += val;
            nd += nd & -nd;
        }
    }

    int query(int nd){ // [1, nd] 구간 쿼리
        int result = 0;
        while (nd > 0){
            result += tree[nd];
            nd -= nd & -nd;
        }
        return result;
    }

    int range_query(int l, int r){ // [l, r] 구간합 계산
        return query(r) - query(l - 1);
    }
}fw;

// heavy-light 분할 (0-based index)
struct HLD{
    int sz[MAXN], pa[MAXN], lv[MAXN], in[MAXN], out[MAXN], head[MAXN], idx;
    vector<int> graph[MAXN];

    void init();

    int _dfs(int here){ // 서브트리의 크기 구하기 및 제일 무거운 서브트리를 첫번째로 보내기
        sz[here] = 1;
        for (int i = 0, gsz = graph[here].size(); i < gsz; i++){
            int there = graph[here][i];
            if (pa[here] == there) continue;
            pa[there] = here;
            lv[there] = lv[here] + 1;
            sz[here] += _dfs(there);
            if (sz[graph[here][0]] < sz[there]) // 무거운 서브트리를 첫번째로
                swap(graph[here][0], graph[here][i]);
        }

        return sz[here];
    }

    void _hld(int here){ // 오일러 경로 구하기 및 체인 설정
        in[here] = ++idx;

        for (auto there: graph[here]){
            if (pa[here] == there) continue;
            if (graph[here][0] == there) // 제일 무거운 서브트리. 즉, 첫번째면 체인을 이어붙이고
                head[there] = head[here];
            else // 아니면 새 체인이 시작된다.
                head[there] = there; 
            _hld(there);
        }

        out[here] = idx;
    }

    void update(int idx, int val){ // 점 업데이트
        fw.update(in[idx] + 1, val);
    }

    int query(int u, int v){ // [u, v] 구간 쿼리
        int result = 0;

        while (head[u] != head[v]){ // 체인이 같지 않다면
            if (lv[head[u]] < lv[head[v]]) // 깊은 체인부터 타고 올라온다.
                swap(u, v);
            result += fw.range_query(in[head[u]] + 1, in[u] + 1);
            u = pa[head[u]];
        }

        if (in[u] > in[v]) // 체인이 같아졌다면 구간합 쿼리를 위해 오일러 경로 순서로 맞추자.
            swap(u, v);

        // 정점이 아닌 간선을 관리하고 있으므로
        // 구간에 포함되지 않는 첫번째 정점에서 나아가는 간선은 포함하지 않는다.
        result += fw.range_query(in[u] + 2, in[v] + 1);

        return result;
    }
}hld;

void FW::init(){
    fill(tree + 1, tree + MAXN + 1, 0);
    // 1번 마을인 Bitburg의 부모 정점은 없다. 그러므로 관리해야 하는 간선이 없다.
    for (int i = 2; i <= n; i++) update(i, 1);
}

void HLD::init(){
    fill(sz, sz + n, 0);
    fill(pa, pa + n, 0);
    fill(lv, lv + n, 0);
    fill(in, in + n, 0);
    fill(out, out + n, 0);
    fill(head, head + n, 0);
    idx = -1;

    _dfs(0); _hld(0);
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n;
    for (int i = 1, a, b; i < n; i++){
        cin >> a >> b;
        hld.graph[--a].push_back(--b);
        hld.graph[b].push_back(a);
    }

    hld.init(); fw.init();

    int m, a, b; char q;
    cin >> m; m += n - 1;
    while (m--){
        cin >> q;
        if (q == 'A'){ // A a b
            cin >> a >> b;
            if (hld.pa[--b] == --a) // 오일러 경로 순서로 맞추자.
                swap(a, b);
            hld.update(a, -1);
        }
        else{ // W a
            cin >> a;
            cout << hld.query(0, --a) << '\n';
        }
    }
}

• Python

import sys; input = sys.stdin.readline
sys.setrecursionlimit(500000)
from math import ceil, log2

# 펜윅트리 (1-based index)
class FW:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (self.n + 1)

    def init(self):
        # 1번 마을인 Bitburg의 부모 정점은 없다. 그러므로 관리해야 하는 간선이 없다.
        for i in range(2, self.n + 1):
            self.update(i, 1)

    def update(self, nd, val): # 점 업데이트
        while nd <= self.n:
            self.tree[nd] += val
            nd += nd & -nd

    def query(self, nd): # [1, nd] 구간 쿼리
        result = 0
        while nd > 0:
            result += self.tree[nd]
            nd -= nd & -nd
        return result

    def range_query(self, l, r): # [l, r] 구간합 계산
        return self.query(r) - self.query(l - 1)

# heavy-light 분할 (0-based index)
class HLD:
    def __init__(self, n):
        self.n = n
        self.graph = [[] for _ in range(self.n)]
        self.sz = [0] * self.n
        self.pa = [0] * self.n
        self.lv = [0] * self.n
        self.inn = [0] * self.n
        self.out = [0] * self.n
        self.head = [0] * self.n
        self.idx = -1

    def init(self):
        self._dfs(0)
        self._hld(0)

    def _dfs(self, here): # 서브트리의 크기 구하기 및 제일 무거운 서브트리를 첫번째로 보내기
        self.sz[here] = 1
        for i in range(len(self.graph[here])):
            there = self.graph[here][i]
            if self.pa[here] == there:
                continue
            self.pa[there] = here
            self.lv[there] = self.lv[here] + 1
            self.sz[here] += self._dfs(there)
            if self.sz[self.graph[here][0]] < self.sz[there]: # 무거운 서브트리를 첫번째로
                self.graph[here][0], self.graph[here][i] = self.graph[here][i], self.graph[here][0]

        return self.sz[here]

    def _hld(self, here): # 오일러 경로 구하기 및 체인 설정
        self.idx += 1
        self.inn[here] = self.idx

        for there in self.graph[here]:
            if self.pa[here] == there:
                continue
            if self.graph[here][0] == there: # 제일 무거운 서브트리. 즉, 첫번째면 체인을 이어붙이고
                self.head[there] = self.head[here]
            else: # 아니면 새 체인이 시작된다.
                self.head[there] = there
            self._hld(there)

        self.out[here] = self.idx

    def update(self, idx, val): # 점 업데이트
        fw.update(self.inn[idx] + 1, val)

    def query(self, u, v): # [u, v] 구간 쿼리
        result = 0

        while self.head[u] != self.head[v]: # 체인이 같지 않다면
            if self.lv[self.head[u]] < self.lv[self.head[v]]: # 깊은 체인부터 타고 올라온다.
                u, v = v, u
            result += fw.range_query(self.inn[self.head[u]] + 1, self.inn[u] + 1)
            u = self.pa[self.head[u]]

        if self.inn[u] > self.inn[v]: # 체인이 같아졌다면 구간합 쿼리를 위해 오일러 경로 순서로 맞추자.
            u, v = v, u

        # 정점이 아닌 간선을 관리하고 있으므로
        # 구간에 포함되지 않는 첫번째 정점에서 나아가는 간선은 포함하지 않는다.
        result += fw.range_query(self.inn[u] + 2, self.inn[v] + 1)

        return result

n = int(input())

hld = HLD(n); fw = FW(n)

for _ in range(n - 1):
    a, b = map(int, input().split())
    a -= 1; b -= 1
    hld.graph[a].append(b)
    hld.graph[b].append(a)

hld.init(); fw.init()

m = int(input()) + n - 1
for _ in range(m):
    query = input().split()
    if query[0] == 'A': # A a b
        a = int(query[1]) - 1; b = int(query[2]) - 1
        if hld.pa[b] == a: # 오일러 경로 순서로 맞추자.
            a, b = b, a
        hld.update(a, -1)
    else: # W a
        a = int(query[1]) - 1
        print(hld.query(0, a))

코드 (수정 전)

• C++

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

int n;
int sz[250000], parent[250000]; // 서브트리의 크기, 부모
int idx = 0, ids[250000]; // 오일러 경로
int head[250000]; // 각 정점이 속해 있는 체인의 가장 위에 있는 정점
int t, tree[1 << (int)ceil(log2(250000) + 1)]; // 펜윅트리
vector<int> graph[250000];

int dfs(int here){ // 서브트리의 크기 구하기 및 제일 무거운 서브트리를 첫번째로 보내기
    sz[here] = 1;
    for (int i = 0, there; i < graph[here].size(); i++){
        there = graph[here][i];
        if (there == parent[here]) continue;
        parent[there] = here;
        sz[here] += dfs(there);
        if (sz[there] > sz[graph[here][0]]) // 무거운 서브트리를 첫번째로
            swap(graph[here][0], graph[here][i]);
    }
    return sz[here];
}

void hld(int here){ // 오일러 경로 구하기 및 체인 설정
    ids[here] = idx++;
    for (int i = 0, there; i < graph[here].size(); i++){
        there = graph[here][i];
        if (there == parent[here]) continue;
        if (!i) head[there] = head[here]; // 제일 무거운 서브트리. 즉, 첫번째면 체인을 이어붙이고
        else head[there] = there; // 아니면 새 체인이 시작된다.
        hld(there);
    }
}

void update(int node, int val){ // 펜윅트리 업데이트
    while (node < n + 1){
        tree[node] += val;
        node += node & -node;
    }
}

int query(int node){ // 펜윅트리 쿼리
    int result = 0;
    while (node > 0){
        result += tree[node];
        node -= node & -node;
    }
    return result;
}

int range_query(int l, int r){ // 펜윅트리 구간합 계산
    return query(r + 1) - query(l); // query(r) - query(l - 1)
}

int hld_query(int u, int v){ // 체인별로 구간합 쿼리
    int result = 0;

    while (head[u] != head[v]){ // 체인이 같지 않다면
        if (sz[head[u]] > sz[head[v]]) swap(u, v); // 무거운(깊은) 체인부터 타고 올라온다.
        result += range_query(ids[head[u]], ids[u]);
        u = parent[head[u]];
    }

    if (ids[u] > ids[v]) swap(u, v); // 체인이 같아졌다면 구간합 쿼리를 위해 오일러 경로 순서로 맞추자.
    result += range_query(ids[u] + 1, ids[v]); // 정점이 아닌 간선을 관리하고 있으므로 구간에 포함되지 않는 첫번째 정점에서 나아가는 간선은 포함하지 않는다.
    return result;
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n;
    int a, b;
    for (int i = 0; i < n - 1; i++){
        cin >> a >> b;
        graph[--a].push_back(--b);
        graph[b].push_back(a);
    }

    // heavy-light 분할
    dfs(0);
    hld(0);

    // 펜윅트리 준비
    // 1번 마을인 Bitburg의 부모 정점은 없다. 그러므로 관리해야 하는 간선이 없다.
    for (int i = 2; i <= n; i++) update(i, 1);

    // 펜윅은 1-based index, 나머지는 0-based index
    int m; char q;
    cin >> m;
    for (int i = 0; i < n - 1 + m; i++){
        cin >> q;
        if (q == 'A'){ // A a b 쿼리
            cin >> a >> b;
            if (parent[--b] == --a) swap(a, b); // 오일러 경로 순서가 b -> a가 되게 맞추자.
            update(ids[a] + 1, -1);
        }
        else{ // W a 쿼리
            cin >> a;
            cout << hld_query(0, --a) << '\n';
        }
    }
}

• Python

import sys; input = sys.stdin.readline
sys.setrecursionlimit(500000)

def dfs(here): # 서브트리의 크기 구하기 및 제일 무거운 서브트리를 첫번째로 보내기
    size[here] = 1
    for i in range(len(graph[here])):
        there = graph[here][i]
        if there == parent[here]:
            continue
        parent[there] = here
        size[here] += dfs(there)
        if size[there] > size[graph[here][0]]: # 무거운 서브트리를 첫번째로
            graph[here][0], graph[here][i] = graph[here][i], graph[here][0]
    return size[here]

def hld(here): # 오일러 경로 구하기 및 체인 설정
    global idx
    ids[here] = idx
    idx += 1
    for i in range(len(graph[here])):
        there = graph[here][i]
        if there == parent[here]:
            continue
        if not i: # 제일 무거운 서브트리. 즉, 첫번째면 체인을 이어붙이고
            head[there] = head[here]
        else: # 아니면 새 체인이 시작된다.
            head[there] = there
        hld(there)

def update(node, val): # 펜윅트리 업데이트
    while node < len(tree):
        tree[node] += val
        node += node & -node

def query(node): # 펜윅트리 쿼리
    result = 0
    while node > 0:
        result += tree[node]
        node -= node & -node
    return result

def range_query(l, r): # 펜윅트리 구간합 계산
    return query(r + 1) - query(l) # query(r) - query(l - 1)

def hld_query(u, v): # 체인별로 구간합 쿼리
    result = 0

    while head[u] != head[v]: # 체인이 같지 않다면
        if size[head[u]] > size[head[v]]: # 무거운(깊은) 체인부터 타고 올라온다.
            u, v = v, u
        result += range_query(ids[head[u]], ids[u])
        u = parent[head[u]]

    if ids[u] > ids[v]: # 체인이 같아졌다면 구간합 쿼리를 위해 오일러 경로 순서로 맞추자.
        u, v = v, u
    result += range_query(ids[u] + 1, ids[v]) # 정점이 아닌 간선을 관리하고 있으므로 구간에 포함되지 않는 첫번째 정점에서 나아가는 간선은 포함하지 않는다.
    return result

n = int(input())
graph = [[] for _ in range(n)]
for _ in range(n - 1):
    a, b = map(int, input().split())
    a -= 1; b -= 1
    graph[a].append(b)
    graph[b].append(a)

# heavy-light 분할
size = [0] * n # 서브트리의 크기
parent = [0] * n # 부모
dfs(0)

ids = [0] * n; idx = 0 # 오일러 경로
head = [0] * n # 각 정점이 속해 있는 체인의 가장 위에 있는 정점
hld(0)

# 펜윅트리 준비
tree = [0] * (n + 1)
for i in range(2, n + 1): # 1번 마을인 Bitburg의 부모 정점은 없다. 그러므로 관리해야 하는 간선이 없다.
    update(i, 1)

# 펜윅은 1-based index, 나머지는 0-based index
m = int(input())
for _ in range(n - 1 + m):
    q = input().split()
    if q[0] == 'A': # A a b 쿼리
        a = int(q[1]) - 1
        b = int(q[2]) - 1
        if parent[b] == a: # 오일러 경로 순서가 b -> a가 되게 맞추자.
            a, b = b, a
        update(ids[a] + 1, -1)
    else: # W a 쿼리
        a = int(q[1]) - 1
        print(hld_query(0, a))