[BOJ 13511] 트리와 쿼리 2

문제 출처 : https://www.acmicpc.net/problem/13511

 

풀이 :

LCA 문제 입니다.

첫번째 쿼리를 처리하기 위해선 아래와 같이 LCA를 구하기 위해 전처리를 할때 cost 값도 같은 점화식으로 구해주시면 되고,

for (int i = 1; depth[ni] > 1 << (i - 1); i++) {
    anc[ni][i] = anc[anc[ni][i - 1]][i - 1];
    cost[ni][i] = cost[anc[ni][i - 1]][i - 1] + cost[ni][i - 1];
}

 

두번째 쿼리를 처리하기 위해선 u와 v의 LCA를 구하고 u,v와 root의 depth를 비교하여 k번째 정점을 찾으시면 됩니다.

 

#include<iostream>
#include<vector>
#include<math.h>
#include<algorithm>
using namespace std;
#define x first
#define y second
typedef pair<int, int> p;
typedef long long ll;
 
int n, m, a, b, c, depth[100001], anc[100001][18], d;
ll cost[100001][18];
vector<p>adj[100001];
 
void dfs(int idx) {
    for (p next : adj[idx]) {
        int nc = next.y, ni = next.x;
        if (depth[ni])continue;
        depth[ni] = depth[idx] + 1;
        anc[ni][0] = idx;
        cost[ni][0] = nc;
        for (int i = 1; depth[ni] > 1 << (i - 1); i++) {
            anc[ni][i] = anc[anc[ni][i - 1]][i - 1];
            cost[ni][i] = cost[anc[ni][i - 1]][i - 1] + cost[ni][i - 1];
        }
        dfs(ni);
    }
}
 
ll query1(int u, int v) {
    ll ret = 0;
    if (depth[u] > depth[v])swap(u, v);
    for (int i = 17; ~i; i--)if (depth[v] - depth[u] >= (1 << i)) {
        ret += cost[v][i];
        v = anc[v][i];
    }
    if (u == v)return ret;
    for (int i = 17; ~i; i--)if (anc[v][i] != anc[u][i]) {
        ret += cost[v][i] + cost[u][i];
        v = anc[v][i];
        u = anc[u][i];
    }
    return ret + cost[v][0] + cost[u][0];
}
 
int lca(int u, int v) {
    if (depth[u] > depth[v])swap(u, v);
    for (int i = 17; ~i; i--)if (depth[v] - depth[u] >= (1 << i))v = anc[v][i];
    if (u == v)return u;
    for (int i = 17; ~i; i--)if (anc[v][i] != anc[u][i]) {
        v = anc[v][i];
        u = anc[u][i];
    }
    return anc[v][0];
}
 
int query2(int u, int v, int k) {
    int root = lca(u, v);
    if (depth[u] - depth[root] + 1 >= k) {
        for (int i = 17; ~i; i--)if ((1 << i) < k) {
            u = anc[u][i];
            k -= (1 << i);
        }
        return u;
    }
    else {
        k = depth[v] - depth[root] - (k + depth[root] - depth[u] - 1);
        if (k == 0)return v;
        for (int i = 17; ~i; i--)if ((1 << i) <= k) {
            v = anc[v][i];
            k -= (1 << i);
        }
        return v;
    }
}
 
int main() {
    ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        cin >> a >> b >> c;
        adj[a].push_back({ b,c });
        adj[b].push_back({ a,c });
    }
    depth[1] = 1;
    dfs(1);
    cin >> m;
    while (m--) {
        cin >> a >> b >> c;
        if (a == 1) cout << query1(b, c) << '\n';
        else {
            cin >> d;
            cout << query2(b, c, d) << '\n';
        }
    }
}