본문 바로가기

ps(알고리즘)

정점들의 거리

이전에 풀었던 lca 알고리즘을 사용하면 되는 문제이다.

https://bluesunset-hack.tistory.com/214

 

LCA 2

제목에서부터 나와있듯 그냥 LCA 알고리즘을 구현하라는 문제이다.정점에서부터 하나씩 부모 노드로 거슬러 올라가게 되면 시간이 너무 많이 걸리기 때문에2의 거듭제곱 단위로 끊어서 거슬러

bluesunset-hack.tistory.com

lca를 사용해서 공통부모를 찾고, 공통부모까지의 거리의 합을 출력하면 끝이다.

#include<stdio.h>
#include<vector>
using namespace std;
int depth[40004];
int N,M;
struct Data
{
    int num, dist;
}parent[40004][30];
vector<Data>nodes[40004];
int height;
void dfs(int n, int par, int dep, int dist)
{
    depth[n] = dep;
    parent[n][0] = {par,dist};
    for(int i=0;i<nodes[n].size();i++)
    {
        if(par!=nodes[n][i].num)
        {
            dfs(nodes[n][i].num,n,dep+1,nodes[n][i].dist);
        }
    }
}

int lca(int a, int b)
{
    if(depth[a] < depth[b])
    {
        int t = a;
        a = b;
        b = t;
    }
    int sum = 0;
    if(depth[a] > depth[b])
    {
        int dif = depth[a] - depth[b];
        int i = 0;
        while(dif>0)
        {
            if(dif%2==1)
            {
                sum += parent[a][i].dist;
                a = parent[a][i].num;
            }
            dif /= 2;
            i++;
        }
    }
    if(a==b) return sum;
    int i = 0;
    for(int k = height;k>=0;k--)
    {
        if(parent[a][k].num!=0 && parent[a][k].num!=parent[b][k].num)
        {
            sum += parent[a][k].dist;
            sum += parent[b][k].dist;
            a = parent[a][k].num;
            b = parent[b][k].num;
        }
    }
    sum += parent[a][0].dist;
    sum += parent[b][0].dist;
    return sum;
}
int main()
{
    scanf("%d",&N);
    for(int i=1;i<=N-1;i++)
    {
        int a,b,c;
        scanf("%d %d %d",&a,&b,&c);
        nodes[a].push_back({b,c});
        nodes[b].push_back({a,c});
    }
    dfs(1,0,0,0);
    height = 0;
    int t = N;
    while(t>1)
    {
        t/=2;
        height++;
    }
    for(int k=1;k<=height;k++)
    {
        for(int i=2;i<=N;i++)
        {
            if(parent[i][k-1].num!=0)
            {
                parent[i][k].num = parent[parent[i][k-1].num][k-1].num;
                parent[i][k].dist = parent[i][k-1].dist + parent[parent[i][k-1].num][k-1].dist;
            }
        }
    }
    scanf("%d",&M);
    for(int i=1;i<=M;i++)
    {
        int a,b;
        scanf("%d %d",&a,&b);
        printf("%d\n",lca(a,b));
    }

}

전체 코드이다.

'ps(알고리즘)' 카테고리의 다른 글

숨바꼭질 5  (0) 2025.02.24
고층 빌딩  (0) 2025.02.24
임계경로  (0) 2025.02.23
사탕상자  (0) 2024.12.04
선분 그룹  (0) 2024.12.03