[C++] 1761번 정점들의 거리 - 트리, 최소 공통 조상, 희소 배열

#include <iostream>
#include <vector>
#include <string.h>
#include <algorithm>

#define FOR(i, k, N) for(int i=k; i<=N; i++)
#define Vr first
#define R second
#define MAX_N 40404

using namespace std;
typedef long long ll;

int N, M;  //노드의 개수 N, 질문의 개수 M
vector<pair<int, int>> Vertex[MAX_N];  //트리 구현. 부모, 자식, 거리
int DT[MAX_N], PA[MAX_N][2];  //Depth, Parent
int vi[MAX_N];  //visited?
ll res;

void input()
{
    cin >> N;
    FOR(i, 1, N-1) {
        int a, b, r;
        cin >> a >> b >> r;
        Vertex[a].push_back({b, r});
        Vertex[b].push_back({a, r});
    }
}

void CD(int no, int Depth)  //깊이 설정하기(lca 용)
{
    vi[no]=1;
    FOR(i, 1, Vertex[no].size()) {
        if(!vi[ Vertex[no][i-1].Vr ]) {
            PA[Vertex[no][i-1].Vr][0]=no;
            PA[Vertex[no][i-1].Vr][1]=Vertex[no][i-1].R;
            CD(Vertex[no][i-1].Vr, Depth+1);
        }
    }
    DT[no] = Depth;
}

void LCA(int a, int b)  //말 그대로 LCA
{
    if(DT[a] < DT[b]) swap(a, b);
    
    while(DT[a]!=DT[b]) {
        res+=PA[a][1]; a = PA[a][0];
    }
    
    while(a!=b) {
        res+=PA[a][1]+PA[b][1];
        a = PA[a][0]; b = PA[b][0];
    }
}

int main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    
    FOR(i, 0, MAX_N) Vertex[i].push_back({0, 0});  //아무것도 들어있지 않으면 .size()에서 오류
    
    input();

    CD(1, 0);  //Calculate Depth
    
    cin >> M;
    FOR(i, 1, M) {
        int a, b;  //x와 y사이의 거리
        cin >> a >> b;
        
        res=0; LCA(a, b);
        cout << res << '\n';
    }
    
    return 0;
}

엄.. 나는 그냥 처음부터 lca를 써서 올라가면서 더해주는 방식으로 해서 시간이 600ms가 소요되었다.

 

다른 사람들 코드를 보니까 28ms로 푼 사람들이 있어서 코드를 봤는데 너무 잘짜여져 있어서 놀라웠다. 참고한 코드가 더 좋은 것 같아서 그 코드의 알고리즘을 아래에 적어두었고, 링크도 달아놨다.

 

 

 

 

 

<참조한 코드(풀어야 볼 수 있습니다.)>

https://www.acmicpc.net/source/7037392

로그인

 

 

<알고리즘>

1. 트리에 저장하는 방식은 같다.
2. 아래까지 내려가면서(dfs) depth를 기록해주고, 거리는 계속 더해준 값을 저장한다.
3. 부모 트리를 저장할 때 2의 배수적으로 저장한다.
4. Depth로 lca를 하면서 올라갈 때 2의 배수적으로 젖아한 부모 트리를 이용하여 올라가며 두 수 a, b의 depth가 같아졌을 때는 1칸씩 올려보낸다.
5. 거리는 a까지의 거리+b까지의 거리 - 공통조상까지의 거리 * 2로 설정한다.

 

 

 

 

 

 

<LCA 알고리즘>

https://kibbomi.tistory.com/201

최소 공통 조상 (LCA, Lowest Common Ancestor)(C/C++)

 

 

 

 

 

 

※현재 고등학교 등교 중인 학생입니다. 이제 알고리즘을 본격적으로 공부하기 시작해서 아직 초보입니다. 혹시 제가 잘못 알고 있는 점이나 더 좋은 풀이 방법이 있어 댓글에 남겨주시면 감사히 하나하나 열심히 읽어보겠습니다. 좋아요, 단순한 댓글 한마디라도 저에겐 큰 힘이 됩니다! 감사합니다.