[BOJ 1761] 정점들의 거리

문제

www.acmicpc.net/problem/1761

1761번: 정점들의 거리

 

---

풀이

LCA(Lowest Common Ancestor)를 사용합니다.

먼저, 다음 4개의 배열을 만듭니다:

 

• \(graph[i][j].x,\ graph[i][j].y:\) \(i\)의 \(j\)번째 자식 노드와 그 노드까지의 거리.
• \(par[i][j]:\) \(i\)의 \(2^j\)번째 조상.
• \(depth[i]:\) \(i\)의 깊이.
• \(dist[i]:\) 루트로부터 \(i\)까지의 거리.

트리를 순회하면서 각 노드마다 위의 정보를 저장해줍니다.

 

다음으로, 입력받은 각 쌍 \((a_{i}, b_{i})\)에 대해 두 노드의 \(lca\)를 구해준 다음,

\(dist[a_{i}]+dist[b_{i}]\)에서 공통부분인 \(dist[lca]\) 부분을 빼주면 답을 알 수 있습니다.

 

---

코드

#include <bits/stdc++.h>
#define x first
#define y second
using namespace std;

typedef long long ll;
typedef pair<int, int> p;
const int MAX_LVL = (int)floor(log2(40000));

int N, M;
vector<vector<p>> graph;
vector<vector<int>> par;
vector<int> depth;
vector<int> dist;

void solve(int cur, int pnode) {
    for (int i = 1; i <= MAX_LVL; ++i)
        par[cur][i] = par[par[cur][i - 1]][i - 1];

    for (auto next : graph[cur]) {
        if (next.x != pnode) {
            depth[next.x] = depth[cur] + 1;
            par[next.x][0] = cur;
            dist[next.x] = dist[cur] + next.y;
            solve(next.x, cur);
        }
    }
}

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

    cin >> N;
    graph.resize(N + 1);
    par.resize(N + 1, vector<int>(MAX_LVL + 1));
    depth.resize(N + 1, 0);
    dist.resize(N + 1);

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

    solve(1, 0);

    cin >> M;
    while (M--) {
        int a, b;
        cin >> a >> b;
        int ans_x = dist[a], ans_y = dist[b];
        if (depth[a] != depth[b]) {
            if (depth[a] >= depth[b])
                swap(a, b);

            for (int i = MAX_LVL; i >= 0; --i)
                if (depth[a] <= depth[par[b][i]])
                    b = par[b][i];
        }

        int lca = a;
        if (a != b) {
            for (int i = MAX_LVL; i >= 0; --i) {
                if (par[a][i] != par[b][i]) {
                    a = par[a][i];
                    b = par[b][i];
                }
                lca = par[a][i];
            }
        }
        cout << ans_x + ans_y - dist[lca] * 2 << "\n";
    }

    return 0;
}