[백준] 13511번 트리와 쿼리 2

13511 트리와 쿼리 2

문제

---

N개의 정점으로 이루어진 트리(무방향 사이클이 없는 연결 그래프)가 있다. 정점은 1번부터 N번까지 번호가 매겨져 있고, 간선은 1번부터 N-1번까지 번호가 매겨져 있다.

아래의 두 쿼리를 수행하는 프로그램을 작성하시오

* 1 u v: u에서 v로 가는 경로의 비용을 출력한다.

* 2 u v k: u에서 v로 가는 경로에 존재하는 정점 중에서 k번째 정점을 출력한다. k는 u에서 v로 가는 경로에 포함된 정점의 수보다 작거나 같다.

입력

---

첫째 줄에 N (2 ≤ N ≤ 100,000)이 주어진다.

둘째 줄부터 N-1개의 줄에는 i번 간선이 연결하는 두 정점 번호 u와 v와 간선의 비용 w가 주어진다.

다음 줄에는 쿼리의 개수 M (1 ≤ M ≤ 100,000)이 주어진다.

다음 M개의 줄에는 쿼리가 한 줄에 하나씩 주어진다.

간선의 비용은 항상 1,000,000보다 작거나 같은 자연수이다.

출력

---

각각의 쿼리의 결과를 순서대로 한 줄에 하나씩 출력한다.

---

풀이)

문제 제목만 보고 세그먼트 트리를 구현하는 문제인가 했는데,

LCA를 이용해서 경로 비용과 경로 상의 k번째 노드를 출력해야하는 문제였다.

1. 경로 비용을 구하는 과정

 - 미리 트리의 루트에서 각 노드까지 이동할 때 드는 비용을 따로 저장해 놓는다.

   long long cost[i] // root ~ i까지 가는 데 드는 비용

  // int형이 아닌 long long으로 선언한 이유는 하나의 간선 비용인 최대 1,000,000이기때문이다.

 - u와 v를 입력받았을 때

   root에서 u까지 가는 비용과 root에서 v까지 가는 비용을 찾아 더한다.

   그리고 그 두 비용에서 root에서 LCA(u,v) 즉, u 정점과 v정점의 최소 공통 부모 정점까지의 비용 * 2을 뺀다.

   비용 = cost[u] + cost[v] - cost[lca(u,v)] * 2 // 두배 곱하는 이유는 각 cost[u], cost[v]에서 겹치는 간선이 존재하기 때문이다.

  그림을 그려보면 이해하기 쉽다.

  예시) 1번 정점이 root인 아래와같은 트리가 있다고 하자(파란색숫자는 간선의 비용을 나타냄)

쿼리문 1 4 5 => 노드 4번에서 5번으로 가는 경로의 비용을 구해보자. 정답 3

lca(4,5) = 2 따라서 비용 = cost[4] + cost[5] - cost[2] * 2 = 2 + 3 - 1 * 2 = 3

2. 경로 상의 k번째 정점 구하는 과정

 - u와 v를 입력받았을때,

   먼저 u에서 u와 v의 최소 공통 부모까지의 정점의 개수를 구한다. 

   // 트리의 깊이차를 이용해 구하며, 최소공통부모까지 포함한 개수를 구한다.

   정점의 개수 = depth[lca(u,v)] - depth[u] + 1

   ① 만일, k와 구한 정점의 개수가 같다면, 최소공통부모 lca(u,v) 정점의 번호를 출력

   ② k가 구한 정점의 개수보다 작다면, u부터 최소공통부모까지 올라가는 경로중에 k번째 정점이 있다는 뜻.

u부터 2^j 꼴로 올라가며 k번째 정점의 번호를 찾아 출력

   ③ k가 구한 정점의 개수보다 크다면, 최소공통부모 다음부터 v까지 내려가는 경로 중에 k번째 정점이 있다는 뜻.

                                            따라서 그 수에 맞춰 v부터 2^j꼴로 올라가며 k번째 정점의 번호를 찾아 출력

 *주의사항 

   경로가 u →...→ lca(u,v) →...→ v임을 잊지 말고 k번째 값을 찾아야한다.

  위 그림의 트리로 다시 예시를 들어보면

  쿼리문 2 4 3 2 => 정점 4번에서 3번으로 가는 경로 중 2번째 정점의 번호를 구해라. 정답 2

  lca(4,3) = 1, 1번 정점부터 4번 정점까지의 경로에서 정점은 총 3개(4,2,1)

  k=2이므로 구한 정점의 개수보다 작다. 따라서 4번부터 위로 올라가면서 경로상의 k번째 찾는다.

  쿼리문 2 6 4 4 => 정점 6번에서 4번으로 가는 경로 중 4번째 정점의 번호를 구해라. 정답 2

  

   lca(6,4) = 1, 1번 정점에서 6번 정점까지의 경로에서 정점은 총 3개(6,3,1)

   k=4이므로 구한 정점의 개수보다 크다. 따라서 4번부터 위로 올라가며 경로상의 k번째 찾는다.

코드)

#include <stdio.h>
#include <vector>
#include <algorithm>
#include <string.h>
using namespace std;
 
#define MAX 18
int n, u, v, w, m;
vector <pair<int,int>> adj[100010];
int depth[100010];
int parent[100010][MAX]; //parent[i][k] = i정점의 2^k번째 부모,log100000 = 16.6...이니깐 넉넉잡아 2^17로
long long cost[100010]; //cost[i] = tree root 정점부터 i번째 정점까지의 비용
 
void dfs(int node) {
    for (int i = 0; i < adj[node].size(); i++) {
        int next = adj[node][i].first;
        if (depth[next] == -1) {
            parent[next][0] = node;
            depth[next] = depth[node] + 1;
            cost[next] += adj[node][i].second + cost[node];
            dfs(next);
        }
    }
}
int main() {
 
    scanf("%d", &n);
    for (int i = 0; i < n - 1; i++) {
        scanf("%d %d %d", &u, &v, &w);
        adj[u].push_back({ v,w });
        adj[v].push_back({ u,w });
    }
    scanf("%d", &m);
    memset(depth, -1, sizeof(depth));
    memset(parent, -1, sizeof(parent));
    
    depth[1] = 0; //트리의 root를 1번 정점으로 
    dfs(1);
 
    for (int k = 0; k < MAX - 1; k++) {
        for (int i = 2; i <= n; i++) {
            parent[i][k + 1] = parent[parent[i][k]][k];
        }
    }
 
    int num, k;
    while (m--)
    {
        scanf("%d %d %d", &num, &u, &v);
        if (num == 1) {//비용출력
            
            long long Temp_c = cost[u] + cost[v];
 
            if (depth[u] < depth[v]) swap(u, v);//u가 더 깊은 정점이 되게
 
            for (int i = MAX - 1; i >= 0; i--) {
                if (depth[u] - depth[v] >= 1 << i) {
                    u = parent[u][i];
                }
            }
            if (u != v) {
                for (int i = MAX - 1; i >= 0; i--) {
                    if (parent[u][i] != -1 && parent[u][i] != parent[v][i]) {
                        u = parent[u][i];
                        v = parent[v][i];
                    }
                }
                u = parent[u][0];
            }
            //root부터 각 정점까지의 비용 - root부터 u,v의 최소공통부모까지의 비용 * 2을 빼기
            //*2를 하는 이유는 root부터 각 정점까지의 비용을 더할때 겹치는 부분이 있으므로
            Temp_c -= 2 * cost[u];
            printf("%lld\n", Temp_c);
 
        }
        else {//k번째 정점 출력
            scanf("%d", &k);
 
            int Temp_u = u;
            int Temp_v = v;
            
            if (depth[u] < depth[v]) swap(u, v); 
 
            for (int i = MAX - 1; i >= 0; i--) {
                if (depth[u] - depth[v] >= 1 << i) {
                    u = parent[u][i];
                }
            }
            if (u != v) {
                for (int i = MAX - 1; i >= 0; i--) {
                    if (parent[u][i] != -1 && parent[u][i] != parent[v][i]) {
                        u = parent[u][i];
                        v = parent[v][i];
                    }
                }
                u = parent[u][0];
            }
            //기존의 u와 v는 Temp_u, Temp_v에 저장
            //u와 v의 공통조상을 u에 저장
            //경로는 Temp_u -> Temp_v로 향해야한다.
            
            int cnt = depth[Temp_u] - depth[u] + 1;//Temp_u에서 공통조상 u까지의 정점개수
 
            if (k == cnt) {
                printf("%d\n", u);
            }
            else if (k < cnt) {
                k--;//시작정점 Temp_u도 개수에 포함
                for (int i = MAX-1; k; i--) {
                    if (k & 1 << i) {
                        Temp_u = parent[Temp_u][i];
                        k -= 1 << i;
                    }
                    
                }
                printf("%d\n", Temp_u);
            }
            else {
                k -= cnt;//최소공통부모 정점에서 아래로 내려가면서 세야할 정점 수
                k = depth[Temp_v]-depth[u]-k; //k개 만큼 아래로 내려야한다면, 반대로 temp_v에서는 몇개 올라와야하는지
                for (int i = MAX - 1; k; i--) {
                    if (k & 1 << i) {
                        Temp_v = parent[Temp_v][i];
                        k -= 1 << i;
                    }
                    
                }
                printf("%d\n", Temp_v);
            }
 
 
        }
    }
    return 0;
}