[PS] 백준 #1761. 정점들의 거리

 

문제 정보

• 난이도: 플래티넘 V
• 문제 주소: https://www.acmicpc.net/problem/1761
• 이 문제를 푼 이유: 알고리즘 스터디 - LCA

 

문제

N(2 ≤ N ≤ 40,000)개의 정점으로 이루어진 트리가 주어지고 M(1 ≤ M ≤ 10,000)개의 두 노드 쌍을 입력받을 때 두 노드 사이의 거리를 출력하라.

 

나의 풀이

두 노드 사이의 거리는 두 노드를 잇는 최단경로의 거리로 정의되는데, 이 최단경로에는 항상 LCA가 포함된다. 따라서 이 문제는 두 노드의 LCA를 찾아서 최단 경로의 거리를 구하는 문제이다. (LCA에 관한 내용은 https://pongkijoa.tistory.com/78?category=543159 에 자세히 서술해두었다.)

이 문제의 핵심은 DFS를 돌 때 각 정점들의 거리 정보를 기록해두어야 한다는 점이다. 이때 거리는 임의의 한 정점에 대해서 구하면 된다. 여기서는 임의로 1번째 노드를 기준으로 거리 정보를 채웠다. 최단 경로는 u -> LCA -> v로 이 거리는 dist[way[0]] + dist[way[2]] - 2*dist[way[1]] 로 구하였다.

#include <bits/stdc++.h>

using namespace std;

#define SIZE 40000
#define MAX 17 // ceil(log2(100000));

int n, m;
int parent[SIZE + 1][MAX + 1];
int depth[SIZE + 1];
vector<vector<pair<int,int>>> adj(SIZE + 1);
vector<int> dist(SIZE + 1);

// dfs로 트리 제작하며 parent[i][0], depth 배열 채움 + dist 채움
void makeTreeByDFS(int curr, int dis) {
    for (auto next : adj[curr]) {
        if (depth[next.first] == -1) {
            parent[next.first][0] = curr;
            depth[next.first] = depth[curr] + 1;
            dist[next.first] = dis + next.second;
            makeTreeByDFS(next.first, dis + next.second);
        }
    }
}

int main() {
    cin.tie(NULL);	cout.tie(NULL);	std::ios::sync_with_stdio(false);

    // 정보 입력
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int u, v, w; cin >> u >> v >> w;
        u--; v--;
        adj[u].push_back(make_pair(v, w)); // 0부터 시작
        adj[v].push_back(make_pair(u, w));
    }

    // 배열 초기화
    memset(parent, -1, sizeof(parent));
    fill(depth, depth + n, -1);
    depth[0] = 0; dist[0] = 0;
    makeTreeByDFS(0, 0);

    // parent 배열 채우기 (parent[i][j]: i번 노드의 2^j 층 위의 조상)
    for (int j = 0; j < MAX; j++) {
        for (int i = 1; i < n; i++) {
            if (parent[i][j] != -1) parent[i][j + 1] = parent[parent[i][j]][j];
        }
    }

    // Finding LCA
    cin >> m;
    while (m--) {
        int u, v; cin >> u >> v;
        u--; v--;
        int way[3]{}; // 경로
        way[0] = u; way[2] = v;

        // 1. 항상 depth[u] >= depth[v]
        if (depth[u] < depth[v]) swap(u, v);
        int diff = depth[u] - depth[v];

        // 2. u와 v가 같은 높이가 되도록 이동
        int j = 0;
        while (diff > 0) {
            if (diff % 2) u = parent[u][j];
            diff /= 2; j++;
        }

        // 3. u와 v가 다르면 2^j 높이 위로 이동 / 같으면 LCA 찾음
        if (u != v) {
            // 높이 2^17, 2^16, ..., 2^2, 2, 1 순으로 시도
            for (int j = MAX - 1; j >= 0; j--) {
                if (parent[u][j] != -1 && parent[u][j] != parent[v][j]) {
                    u = parent[u][j];
                    v = parent[v][j];
                }
            }
            // 두 정점 u, v의 부모가 같으므로 한 번 더 올림
            u = parent[u][0];
        }
        way[1] = u; 
        
        // 4. way 에 따라서 dist 값 구하기 (way[1]: LCA)
        // cout << "LCA 경로: " << way[0] << " -> " << way[1] << " -> " << way[2] << endl;
        int ans = dist[way[0]] + dist[way[2]] - 2*dist[way[1]];
        cout << ans << "\n";
    }
}