LCA 2

제목에서부터 나와있듯 그냥 LCA 알고리즘을 구현하라는 문제이다.

정점에서부터 하나씩 부모 노드로 거슬러 올라가게 되면 시간이 너무 많이 걸리기 때문에

2의 거듭제곱 단위로 끊어서 거슬러 올라가는 방식이다. 이렇게 하면 logN의 시간이 걸려서 훨씬 빠르다.

 

우선 두 정점의 깊이를 맞춰주고, 동시에 거슬러 올라가면 된다.

parent 배열이 중요한데 parent[i][j]가 뜻하는건 i번 정점의 2^j 번째 부모 노드이다.

void set_par_depth(int par,int n,int dep)
{
    depth[n]=dep;
    parent[n][0]=par;
    for(int i=0;i<nodes[n].size();i++)
    {
        if(nodes[n][i]!=par) set_par_depth(n,nodes[n][i],dep+1);
    }
}

우선 depth배열을 채우고 parent[i][0]을 채우는 함수이다.

단순한 dfs 함수이다.

set_par_depth(0,1,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]!=0) parent[i][k]=parent[parent[i][k-1]][k-1];
    }
}

그리고 이렇게 트리의 전체 깊이와 나머지 parent 배열들을 dp 느낌으로 채워준다.

int lca(int a,int b)
{
    if(depth[a]<depth[b])
    {
        int t = a;
        a=b;
        b=t;
    }
    if(depth[a]>depth[b])
    {
        int dif = depth[a]-depth[b];
        int i=0;
        while(dif>0)
        {
            if(dif%2==1) a=parent[a][i];
            dif/=2;
            i++;
        }
    }
    if(a==b) return a;
    int i=0;
    for(int k=height;k>=0;k--)
    {
        if(parent[a][k]!=0 && parent[a][k]!=parent[b][k])
        {
            a=parent[a][k];
            b=parent[b][k];
        }
    }
    return a = parent[a][0];
}

그리고 이렇게 lca 함수에서 depth를 맞춰주고, 동시에 거슬러 올라가면서 공통부모를 찾아준다.

#include<stdio.h>
#include<vector>
using namespace std;
int N,M;
int parent[100004][104];
int depth[100004];
int height=0;
vector<int>nodes[100004];
int lca(int a,int b)
{
    if(depth[a]<depth[b])
    {
        int t = a;
        a=b;
        b=t;
    }
    if(depth[a]>depth[b])
    {
        int dif = depth[a]-depth[b];
        int i=0;
        while(dif>0)
        {
            if(dif%2==1) a=parent[a][i];
            dif/=2;
            i++;
        }
    }
    if(a==b) return a;
    int i=0;
    for(int k=height;k>=0;k--)
    {
        if(parent[a][k]!=0 && parent[a][k]!=parent[b][k])
        {
            a=parent[a][k];
            b=parent[b][k];
        }
    }
    return a = parent[a][0];
}
void set_par_depth(int par,int n,int dep)
{
    depth[n]=dep;
    parent[n][0]=par;
    for(int i=0;i<nodes[n].size();i++)
    {
        if(nodes[n][i]!=par) set_par_depth(n,nodes[n][i],dep+1);
    }
}
int main()
{
    scanf("%d",&N);
    for(int i=1;i<=N-1;i++)
    {
        int a,b;
        scanf("%d %d",&a,&b);
        nodes[a].push_back(b);
        nodes[b].push_back(a);
    }
    set_par_depth(0,1,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]!=0) parent[i][k]=parent[parent[i][k-1]][k-1];
        }
    }
    scanf("%d",&M);
    for(int i=1;i<=M;i++)
    {
        int a,b;
        scanf("%d %d",&a,&b);
        printf("%d\n",lca(a,b));
    }
}

전체 코드이다.