[BOJ알고리즘, C++]#13511_트리와 쿼리2, 트리, LCA, 노드 사이 거리

 

#1. 문제

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

 

---

 

#2. 풀이

 

1. 트리

[자료구조]#5_트리

트리 자료구조는 1:N 관계의 계층 구조를 갖는 그래프의 한 종류입니다. 트리 자료구조의 각 노드 간 연결 관계는 간선으로 나타내며, 순환 구조를 갖지 않기 때문에 "계층 구조"를 형성합니다. 트리 자료구조의 노드 개수가 N개라면, 간선의 개수는 N-1개입니다. 

 

2. LCA

[알고리즘]#7_LCA(Least Common Ancestor)

LCA는 트리 자료구조에서 두 노드의 가장 가까운 조상 노드를 찾는 알고리즘입니다. 이때, 가장 가까운 조상 노드란, 두 노드 사이의 경로에 위치한 노드 중 깊이가 가장 깊은 노드를 의미합니다. LCA 알고리즘은 먼저, 트리 자료구조의 각 노드의 깊이 값, 부모 노드 정보를 DFS를 통해 파악하고, 이를 통해 LCA 알고리즘을 활용합니다.

 

3. u -> lca(u,v) -> v 경로 사이에 k번째 노드 찾기

@LCA를 활용한 두 노드 사이의 거리 값
dist(u,v) = dist(root, u) + dist(root, v) - 2*dist(root,lca(u,v))

@LCA를 활용한 두 노드 사이의 경로 상에서 K번째 노드 찾기
1. lca(u,v)를 통해 u 노드와 v 노드의 가장 가까운 공통 조상 찾기
2. u -> lca(u,v) 사이의 경로와 lca(u,v) -> v 사이의 경로 중 K번째 정점은 어디에 위치하는지 확인
3. 만약, u -> lca(u,v) 경로에 위치한다면, K번째 노드는 u로부터 2^k 번째에 위치한 노드.
4. 만약, lca(u,v) -> u 경로에 위치한다면, K번째 노드는 v로부터 2^(dist(u, lca(u,v)) + 1 + (depth[v] - depth[lac(u,v)]) - k) 번째에 위치한 노드

1. 쿼리 1번은 간단합니다. LCA를 활용해 두 노드 사이의 거리 값을 찾습니다.
2. 쿼리 2번은 u -> v 경로 상에 K번째 위치한 노드를 찾는 문제입니다. u -> v 경로를 u -> lca(u, v) -> v 경로로 가정하고, u -> lca(u, v) 경로 상에 위치하는지, lca(u, v) -> v 경로 상에 위치하는지 확인하고, 각각 u로부터 2^k번째 부모 노드가 정답이 되고, v로부터 위 공식을 통해 올라가 부모 노드가 정답이 됩니다.

 

---

 

#3. 코드

/*
    @문제 : 트리 자료구조에서 두 노드사이의 거리, 두 노드 사이의 경로에서 k번째 정점 찾기
    @설명
        1. LCA 활용
        2. 쿼리1 : dist(u,v) = dist(root,u) + dist(root,v) - 2*(dist(root, lca(u,v)))
        3. 쿼리2 : u -> lca(u,v) -> v 경로에서 u로부터 k번째 정점을 찾기
                1. 먼저, k가 u -> lca(u,v) 경로에 있는지, 혹은 lca(u,v) -> v 경로에 있는지 확인합니다.
                    u -> lca(u,v) 거리가 k보다 크다면, k는 u -> lca(u,v) 경로에 있습니다!
                2. 만약, k가 u -> lca(u,v) 경로에 있다면, u로부터 2^k
*/

#include <iostream>
#include <vector>
using namespace std;

#define MAX 100001
#define LOG 17

typedef long long ll;

int N, M;
ll parent[MAX][LOG];
ll depth[MAX];
ll dist[MAX];
vector<pair<int,ll>> tree[MAX];

// #1. LCA를 위한 DFS(깊이 목록 + 부모 목록 초기화)
void dfs(ll cur, ll p, ll d)
{
    depth[cur] = depth[p]+1;
    dist[cur] = d;
    parent[cur][0] = p;
    // #1. Binary Lifting, 부모 목록 초기화
    for(int k=1; k<LOG; ++k)
    {
        if(parent[cur][k-1] != -1)
        {
            parent[cur][k] = parent[parent[cur][k-1]][k-1];
        }
        else parent[cur][k] = -1;
    }
    // #2. DFS
    for(const auto edge : tree[cur])
    {
        int neighbor = edge.first;
        ll weight = edge.second;

        if(neighbor == p) continue;

        dfs(neighbor, cur, dist[cur]+weight);
    }
}

// #2. LCA
int lca(int u, int v)
{
    // u > v 유지
    if(depth[u] < depth[v]) swap(u,v);

    // 깊이 맞추기
    for(int k=LOG-1; k>=0; --k)
    {
        if(depth[u]-depth[v] >= (1<<k)) u = parent[u][k];
    }

    // u == v?
    if(u==v) return u;

    // lca 찾기
    for(int k=LOG-1; k>=0; --k)
    {
        if(parent[u][k] != parent[v][k])
        {
            u = parent[u][k];
            v = parent[v][k];
        }
    }

    // 결과 반환
    return parent[u][0];
}

// #3. 무 방향에 가중치가 존재하기 때문에, LCA를 활용
ll Dist(int start, int end, int lcaVal)
{
    // LCA를 통해 거리 구하기 공식 : dist(root, u) + dist(root, v) - 2*dist(root, lca(u,v))
    return dist[start] + dist[end] - 2*dist[lcaVal];
}

// #4. 현재 노드로부터 2^k번째 위에 위치한 부모 노드
// *중요: 현재 위치로부터 1칸 위로 이동, parent[node][0] -> parent[node][1]로 이동하는 작업은 1만큼 이동하는 것이 아니라, 2^1 만큼 이동한다.
//        따라서, k가 5라면 이진수로 101이며, 2^0 + 2^2 가 되며, i가 0, 그리고 2일 때만 위로 움직여주는 작업을 수행해주는겁니다. 덤으로 k를 오른쪽으로 1칸 시프트 해주며!
int moveUp(int node, int k) {
    for (int i = 0; k > 0 && i < LOG; ++i) {
        if (k & 1) node = parent[node][i]; // k를 binary로 표현했을 때, 마지막 비트가 1인지 확인(홀수인지 확인!)
        k >>= 1;
    }
    return node;
}

// #5. u -> lca(u,v) -> v 경로 상에서 k번째 위치한 정점 찾기
int KthNode(int start, int end, int k, int lcaVal)
{
    int distToLCA = depth[start] - depth[lcaVal];
    
    if (k <= distToLCA + 1) {
        // k가 u -> lca(u,v) 경로 상에 위치할 경우
        return moveUp(start, k - 1);
    } else {
        // k가 lca(u,v) -> v 경로 상에 위치할 경우
        int distToEnd = depth[end] - depth[lcaVal];
        int totalDist = distToLCA + 1 + distToEnd; // u -> lca(u,v) -> v 전체 거리
        int kFromEnd = totalDist - k + 1; // v에서 시작하여 k번째 위치
        return moveUp(end, kFromEnd - 1);
    }
}

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

    cin >> N;

    for(int i=1; i<=N; ++i)
        for(int j=0; j<LOG; ++j)
            parent[i][j] = -1;

    depth[0] = -1;

    // #1. 트리 초기화(무방향, 무순환 구조)
    for(int i=0; i<N-1; ++i)
    {
        int u, v, w;
        cin >> u >> v >> w;
        // 무방향! 두 노드에 대하여 서로 출발점과 도착점 모두 될 수 있다!
        tree[u].push_back({v,w});
        tree[v].push_back({u,w});
    }

    // #2. dfs
    dfs(1,0,0);

    // #3. 쿼리
    cin >> M;

    while(M--)
    {
        int query, u, v, k;
        cin >> query;

        if(query == 1)
        {
            cin >> u >> v;
            int res = lca(u,v);
            cout << Dist(u,v, res) << '\n';
        }
        else if(query == 2)
        {
            cin >> u >> v >> k;
            int res = lca(u,v);
            cout << KthNode(u, v, k, res) << '\n';
        }
    }

    return 0;
}

 

---