트리와 쿼리 2 (백준 13511)

안녕하세요. 오늘은 트리와 쿼리 문제를 풀어볼 거예요.

문제

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

아이디어

이 문제는 그냥 함수를 여러개 만들면 됩니다.
LCA(x,y): x와 y의 최소 공통 조상을 반환
dist(x,y): y가 x의 조상임. x와 y의 거리
GetCost(x,y): y가 x의 조상임. x와 y사이의 비용
Kth(x,k): x의 k번째 조상
DFS(node,up): LCA를 만드는 DFS함수
또한 k번째 노드를 찾는것 또한 중요한데 자기자신이 첫번째 노드이고 k번째 노드가 lca를 기준으로 어디에 있는지를 잘 생각해야합니다.

소스코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
#define ll long long
using namespace std;

ll parent[101010][20] = { 0 }, depth[101010] = { 0 }, cost[101010][20] = {0};
vector <pair <ll, ll> > v[101010];

ll LCA(ll x, ll y) //x,y의 최소공통조상 찾기
{
    if (depth[x] < depth[y]) swap(x, y);
    for (ll i = 19; i >= 0; i--)
    {
        if (depth[parent[x][i]] >= depth[y])
            x = parent[x][i];
    }
    if (x == y) return x;
    for (ll i = 19; i >= 0; i--)
    {
        if (parent[x][i] != parent[y][i])
        {
            x = parent[x][i];
            y = parent[y][i];
        }
    }
    return parent[x][0];
}
ll dist(ll x, ll y) //y가 x의 조상이라는 전제하에 두 노드의 거리
{
    ll ans = 0;
    for (ll i = 19; i >= 0; i--)
    {
        if (depth[parent[x][i]] >= depth[y])
        {
            ans += pow(2, i);
            x = parent[x][i];
        }
    }
    return ans;
}
ll GetCost(ll x, ll y) //y가 x의 조상이라는 전제하에 두 노드사이의 비용
{
    ll ans = 0;
    for (ll i = 19; i >= 0; i--)
    {
        if (depth[parent[x][i]] >= depth[y])
        {
            ans += cost[x][i];
            x = parent[x][i];
        }
    }
    return ans;
}
ll Kth(ll x, ll k) //x의 k번 조상
{
    for (ll i = 19; i >= 0; i--)
    {
        if (k >= pow(2, i))
        {
            k -= pow(2, i);
            x = parent[x][i];
        }
    }
    return x;
}

void DFS(ll node, ll up)
{
    parent[node][0] = up;
    depth[node] = depth[up] + 1;
    for (ll i = 1; i <= 19; i++)
    {
        parent[node][i] = parent[parent[node][i - 1]][i - 1];
        cost[node][i] = cost[node][i - 1] + cost[parent[node][i - 1]][i - 1];
    }
    for (auto i : v[node])
        if (i.first != up)
        {
            cost[i.first][0] = i.second;
            DFS(i.first, node);
        }

}

int main(void)
{
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    ll N, i, a, b, c;
    cin >> N;
    for (i = 0; i < N - 1; i++)
    {
        cin >> a >> b >> c;
        v[a].push_back({ b,c });
        v[b].push_back({ a,c });
    }
    DFS(1, 0);
    
    ll M, type, x, y, z;
    cin >> M;
    for (i = 0; i < M; i++)
    {
        cin >> type;
        if (type == 1)
        {
            cin >> x >> y;
            ll lca = LCA(x, y);
            cout << GetCost(x, lca) + GetCost(y, lca) << "\n";
        }
        else
        {
            cin >> x >> y >> z; z--;
            ll lca = LCA(x, y);
            ll dis = dist(x, lca);
            if (z <= dis) //x-LCA의 경로에 k번째 노드가 있는 경우
            {
                cout << Kth(x, z) << "\n";
            }
            else
            {
                cout << Kth(y, dist(y, lca) - (z - dis)) << "\n";
            }
        }
    }
}

감사합니다.