정점들의 거리 - 1761

//1761 - 정점들의 거리
 
#include <iostream>
#include <vector>
using namespace std;
 
int N, M;
vector<vector<pair<int, int>>> adj;
int parent[40001][17];
int cost[40001][17];
int dept[40001];
 
void dfs(int here)
{
    for(int i = 0; i<adj[here].size(); ++i)
    {
        int there = adj[here][i].first;
        int c = adj[here][i].second;
        if(parent[there][0] == 0)
        {
            parent[there][0] = here;
            cost[there][0] = c;
            dept[there] = dept[here] + 1;
            dfs(there);
        }
    }
}
 
void init()
{
    cin >> N;
    adj = vector<vector<pair<int, int>>>(N+1);
    int x, y, r;
    for(int i = 1; i<N; ++i)
    {
        cin >> x >> y >> r;
        adj[x].push_back(make_pair(y, r));
        adj[y].push_back(make_pair(x, r));
    }
 
    parent[1][0] = 1;
    dfs(1);
 
    for(int j = 1; j<17; ++j)
        for(int i = 1; i<=N; ++i)
        {
            parent[i][j] = parent[parent[i][j-1]][j-1];
            cost[i][j] = cost[i][j-1] + cost[parent[i][j-1]][j-1];
        }
}
 
int LCA()
{
    int ret = 0;
    int u, v;
    cin >> u >> v;
    if(dept[u] < dept[v]) swap(u, v);
 
    while(dept[u] != dept[v])
    {
        for(int i = 16; i>=0; --i)
            if(dept[parent[u][i]] > dept[v])
            {
                ret += cost[u][i];
                u = parent[u][i];
                break;
            }
        if(dept[parent[u][0]] == dept[v])
        {
            ret += cost[u][0];
            u = parent[u][0];
            break;
        }
    }
 
    if(u == v) return ret;
 
    while(true)
    {
        for(int i = 16; i>=0; --i)
        {
            if(parent[u][i] != parent[v][i])
            {
                ret += cost[u][i] + cost[v][i];
                u = parent[u][i], v = parent[v][i];
                break;
            }
        }
        if(parent[u][0] == parent[v][0])
            break;
    }
    return ret + cost[u][0] + cost[v][0];
}
 
int main()
{
    ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    init();
    cin >> M;
    for(int i = 0; i<M; ++i)
        cout << LCA() << "\n";
}

 

 

 

 

 

 

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

1761번: 정점들의 거리

 

 

 

 

 

 

정점들의 거리

LCA를 사용해 두점 사이의 거리를 구하면 됩니다.

처음에 TLE가 났는데 59번째 줄에 while(true)로 작성했어서 애초에 dept이 같은 경우에 대해 dept을 맞춰주는 반복문에서 무한루프가 걸려서 나는것이였습니다.