BOJ 1761

정점들의 거리

문제출처

문제 해석

  • N개의 정점으로 이루어진 트리가 주어지고, M개의 두 노드 쌍을 입력받을 때 두 노드 사이의 거리를 출력하는 문제입니다.

설계(20)

  • 앞서 구현한 LCA를 약간 응용한 문제입니다.
  • 기존과 다른 점은 정점간 거리가 다르다는 것인데, DFS탐색을 진행할 때 해당 정점까지의 누적거리를 계산하도록 합니다.
  • 누적거리를 알기 때문에 두 정점의 누적거리를 더하고 중복되어있는 공통조상까지의 거리를 빼주면 두 정점의 거리를 구할 수 있게 됩니다.
  • 수식으로 보았을 때, D[a,b] = dist[a] + dist[b] - 2*dist[lca(a,b)] 와 같이 나옵니다.

구현(60)

코드
#include <iostream>
#include <vector>
#include <cmath>
#define MAXN 40000+1
#define swap(a,b){int t = a; a= b; b = t;}

using namespace std;

vector<pair<int,int> > adj[MAXN];
int depth[MAXN];
int dist[MAXN];
int par[MAXN][20];
bool visited[MAXN];
int max_level = 20;
int N,M;

void getTree(int idx, int d, int cnt){
    visited[idx] = true;
    depth[idx] = d;
    dist[idx] = cnt;
    for (int i = 0 ; i <adj[idx].size(); i++){
        int there = adj[idx][i].first;
        if (visited[there]) continue;
        getTree(there,d+1,cnt+adj[idx][i].second);
        par[there][0] = idx;
    }
}
int LCA(int a, int b){
    // a와 b의 depth 같게 만들기
    if (depth[a] != depth[b]){
        if (depth[a] > depth[b]) swap(a,b);
        for (int i = max_level-1; i >=0; i--){
            if (depth[a] <= depth[par[b][i]]){
                b = par[b][i];
            }
        }
    }
    int lca = a;
    if (a!=b){
        for (int i = max_level-1; i>=0; i--){
            if (par[a][i]!= par[b][i]){
                a = par[a][i];
                b = par[b][i];
            }
            lca = par[a][i];
        }
    }
    return lca;
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> N;
    for (int i = 0 ;i <N-1;i++){
        int u,v,w;
        cin >> u >> v >> w;
        adj[u].push_back({v,w});
        adj[v].push_back({u,w});
    }
    // 트리 만들기
    getTree(1,0,0);
    for (int i = 1; i < max_level; i++){
        for (int j = 1; j <= N; j++){
            par[j][i] = par[par[j][i-1]][i-1];
        }
    }
    cin >> M;
    while (M--){
        int a,b;
        cin >> a >> b;
        cout << dist[a] + dist[b] - 2*dist[LCA(a,b)]<<"\n";
    }
    return 0;
}
디버깅
  • 아직 lca를 직접 구현하는게 익숙치 않아서 구현하는데 시간이 오래걸렸습니다.
제출결과
  • 36ms

마무리

  • lca를 이용하여 두 정점간 거리를 구하는 문제였습니다. O(MlogN)으로 두 정점간 거리를 구할 수 있어 알고리즘 최적화에 많이 쓰일 것 같습니다.