백준 13511번 C++ 풀이

문제

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

시간 제한	메모리 제한	solved.ac 티어
2초	512MB	플래티넘 3

사진을 클릭하면 문제로 이동합니다.

---

풀이

 이 문제는 처음에 희소 테이블을 이용하는 문제인 줄 알았지만, 해보다보니 LCA를 응용해야하는 문제임을 알았다. (태그 까보고...)

 

 임의의 두 노드 u, v에서 u -> v로 가는 경로라 함은 암묵적으로 최단경로를 의미한다. 그렇다면 u -> v의 경로는 LCA를 반드시 거칠 것이다. LCA를 구하기 위한 희소 테이블에 해당 조상으로 가기 위한 경로의 비용도 같이 저장한다. 이렇게 하면 LCA와 비용 모두 빠르게 구할 수 있다.

 

 쿼리를 분석해보자. 1번 쿼리는 u -> v로 가는 경로의 비용을 출력하는 쿼리이다. u -> v의 경로는, u -> LCA_uv -> v의 경로와 같다. 우리는 u → LCA의 비용과 v → LCA의 비용을 LCA를 구하면서 비용도 같이 저장했기 때문에, 이 둘의 합을 출력하면 된다.

 

 2번 쿼리는 u → v의 경로상에서 k번째의 정점을 출력하는 것이다. 이때, 예제입력을 보면 알겠지만, k는 1부터 시작한다. u → v로 가기 위한 정점이 u → a1 → a2 → a3 → ... → LCA → b1 → b2 → b3 → ... → v라고 하고, 정점의 개수가 p개라고 하면, k = 1일땐 u, k = p일땐 v이다. k번째 정점은 a 구간에 있거나(u → LCA), b 구간에 있을 것이다. LCA일수도 있다. 그래서 u → LCA의 정점 개수와 v → LCA의 정점의 개수를 구하고, 경우를 나누어 구해주면 된다.

코드

#include <bits/stdc++.h>
#define endl '\n'
using namespace std;
using ll = long long;
#define MAX_A 20

int n,m;
vector<vector<pair<int,ll>>> tree, graph;
vector<int> depths;
vector<tuple<int,int,int,int>> query;

void make_tree(int st, int p, int c, int h) {
    tree[st][0] = {p, c}, depths[st] = h;

    for(auto& [g, cc] : graph[st]) {
        if (g == p) continue;
        make_tree(g, st, cc, h + 1);
    }
}

void make_ancestor() {
    for (int ancestor = 1; ancestor <= MAX_A; ancestor++) {
        for (int node = 1; node <= n; node++) {
            tree[node][ancestor].first = tree[tree[node][ancestor - 1].first][ancestor - 1].first;
            tree[node][ancestor].second = tree[tree[node][ancestor - 1].first][ancestor - 1].second + tree[node][ancestor - 1].second;
        }
    }
}

void input(){
    int a, b, c;
    cin >> n;
    graph.resize(n + 1);
    tree.resize(n + 1, vector<pair<int,ll>>(MAX_A + 1));
    depths.resize(n + 1, -1);
    for (int i = 1; i < n; i++){
        cin >> a >> b >> c;
        graph[a].emplace_back(b, c);
        graph[b].emplace_back(a, c);
    }
    cin >> m;
    query.resize(m);
    for(auto& [q, u, v, k] : query) {
        cin >> q >> u >> v;
        if (q == 2) cin >> k;
    }
    make_tree(1, 0, 0, 1);
    make_ancestor();
}

// lca, cost, u->lca, v->lca
tuple<int, ll, int, int> get_lca(int u, int v) {
    bool swap_flag = false;
    tuple<int, ll, int, int> ans;
    auto& [lca, cost, du, dv] = ans;
    du = 1, dv = 1;

    if (depths[u] < depths[v]){
        swap(u, v);
        swap_flag = true;
    }

    int dh = depths[u] - depths[v];
    swap_flag ? dv += dh : du += dh;

    for (int i = MAX_A; i >= 0 && dh; i--) {
        if (dh >= (1 << i)) {
            cost += 1LL * tree[u][i].second;
            u = tree[u][i].first;
            dh -= (1 << i);
        }
    }

    if (u != v) {
        for (int i = MAX_A; i >= 0; i--) {
            if (tree[u][i].first != 0 && tree[u][i].first != tree[v][i].first) {
                cost += 1LL * tree[u][i].second + tree[v][i].second;
                u = tree[u][i].first;
                v = tree[v][i].first;
                du += 1 << i, dv += 1 << i;
            }
        }
        cost += 1LL * tree[u][0].second + tree[v][0].second;
        ++du, ++dv;
    }

    lca = tree[v][0].first;

    return ans;
}

void sol() {
    for(auto& [q, target, goal, k] : query) {
        auto [lca, cost, dt, dg] = get_lca(target, goal);
        
        if (q == 1) {
            cout << cost;
        }
        else {
            if (dt >= k) {
                --k;
                for(int i = MAX_A; i >= 0 && k; i--) {
                    if (k >= 1 << i) {
                        target = tree[target][i].first;
                        k -= 1 << i;
                    }
                }
                cout << target;
            }
            else {
                k = dg - (k - dt) - 1;
                for(int i = MAX_A; i >= 0 && k; i--) {
                    if (k >= 1 << i) {
                        goal = tree[goal][i].first;
                        k -= 1 << i;
                    }
                }
                cout << goal;
            }
        }
        cout << endl;
    }
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    input();
    sol();
}