[C++ 백준 풀이][Platinum] 1761번 : 정점들의 거리 / 8012번 : 한동이는 영업사원! / 7742번 : Railway

이 포스트에서는 프로그래밍 문제 사이트 백준 Online Judge의 1761번 : '정점들의 거리' 문제에 대한 풀이 코드와 해설을 다루고 있습니다.

 

문제의 난이도는 Solved.ac 기준 Platinum 티어에 해당하며, 문제를 풀이하기 위해 LCA(최소 공통 조상) 알고리즘에 대한 이해가 필요합니다.

 

 

1761번 : 정점들의 거리

 

트리가 주어지고, 두 노드 쌍의 번호가 입력되면 그 거리를 출력하는 문제입니다.

문제 자체는 아주 간단하지만, N이 4만 이하이고 M이 1만 이하이기 때문에 효율적인 알고리즘을 활용하지 않으면 시간 초과에 걸리게 됩니다.

 

[C++ 백준 풀이] LCA, 최소 공통 조상 알고리즘 (11438번 : LCA 2 / 11437번 : LCA)

 

그러나 이전에 효율적인 LCA 알고리즘에 대해서 다룬 적이 있기 때문에, 이를 활용하여 문제를 풀이해보겠습니다.

(log H 시간에 탐색이 가능한 LCA 알고리즘에 대해 모르시는 분들은 위의 게시글을 참고해주시면 감사하겠습니다.)

 

 

#include <cstdio>
#include <vector>
#define MAX 100001
using namespace std;
 
int Parent[MAX][20], Dist[MAX][20] = {}, Depth[MAX], H = 0;
vector<pair<int, int>> Line[MAX];
 
void setParent(int Par, int Node, int Dis, int Dep) {
    if(!Line[Node].size()) return;
    Parent[Node][0] = Par;
    Dist[Node][0] = Dis;
    Depth[Node] = Dep;
    for(int i=0; i<Line[Node].size(); i++)
        if(Line[Node][i].first != Par) setParent(Node, Line[Node][i].first, Line[Node][i].second, Dep+1);
}
 
int findLCADis(int A, int B) {
    int Dis = 0;
    if(Depth[A] != Depth[B]) {
        if(Depth[A] < Depth[B]) swap(A, B);
        int Diff = Depth[A] - Depth[B];
        for(int i=0; Diff>0; i++) {
            if(Diff%2) {
                Dis += Dist[A][i];
                A = Parent[A][i];
            }
            Diff >>= 1;
        }
    }
    if(A != B) {
        for(int i=H; i>=0; i--)
            if(Parent[A][i] != 0 && Parent[A][i] != Parent[B][i]) {
                Dis += (Dist[A][i] + Dist[B][i]);
                A = Parent[A][i];
                B = Parent[B][i];
            }
        Dis += (Dist[A][0] + Dist[B][0]);
        A = Parent[A][0];
    }
    return Dis;
}
 
int main() {
    int N, M, A, B, C, temp;
    scanf("%d", &N);
    temp = N;
    while(temp > 1) {
        temp /= 2;
        H++;
    }
    for(int i=0; i<N-1; i++) {
        scanf("%d %d %d", &A, &B, &C);
        Line[A].push_back({B, C});
        Line[B].push_back({A, C});
    }
    setParent(0, 1, 0, 0);
    for(int i=1; i<=H; i++)
        for(int j=2; j<=N; j++)
            if(Parent[j][i-1]) {
                Parent[j][i] = Parent[Parent[j][i-1]][i-1];
                Dist[j][i] = Dist[j][i-1] + Dist[Parent[j][i-1]][i-1];
            }
    scanf("%d", &M);
    for(int i=0; i<M; i++) {
        scanf("%d %d", &A, &B);
        printf("%d\n", findLCADis(A, B));
    }
}
 

 

풀이는 기존의 LCA 알고리즘과 거의 비슷한데, 저는 여기에 Dist 배열까지 추가해서 풀었습니다.

먼 거리의 부모 노드를 이동할 때 2의 거듭제곱의 조합을 이용해서 Parent들을 건너뛰는 방식으로 알고리즘이 사용되었는데, 거리를 저장하는 Dist 배열 역시 같은 방식으로 갱신하여, 먼 거리의 Parent까지의 거리의 합을 log H 시간에 구할 수 있도록 구현하였습니다.

핵심이 되는 식은 62행의 Dist[j][i] = Dist[j][i-1] + Dist[Parent[j][i-1]][i-1]입니다.

기존 LCA 알고리즘에 대한 이해가 있어야 거리가 올바르게 갱신되도록 알맞은 식을 작성할 수 있습니다.

 

나머지는 그냥 findLCA 함수를 적절히 손봐서 Dist 값의 합을 계산해주기만 하면 됩니다.

 

 

 

풀이를 제출해보면 44ms의 짧은 시간에 모든 테스트케이스를 통과했음을 확인할 수 있습니다.

 

 

8012번 : 한동이는 영업사원!

 

1번 노드에서 시작하여, 주어진 순서대로 모든 노드를 방문한다고 할 때 몇 개의 간선을 이동해야 하는지를 구하는 문제입니다. (모든 도시 사이의 이동에는 1의 시간이 걸린다고 하였으므로)

이는 위에서 풀이한 트리에서의 두 정점 사이의 거리를 응용하여, M개의 순서가 있을 때 각 인접한 순서 사이의 거리, 즉 M-1개의 거리의 합을 구해주면 됩니다.

 

이 때 모든 거리는 1이라고 설정하면 되기 때문에 매개 변수 등의 설정이 더 간단해질 수 있습니다.

 

 

#include <cstdio>
#include <vector>
#define MAX 100001
using namespace std;
 
int Parent[MAX][20], Dist[MAX][20] = {}, Depth[MAX], H = 0;
vector<int> Line[MAX];
 
void setParent(int Par, int Node, int Dep) {
    if(!Line[Node].size()) return;
    Parent[Node][0] = Par;
    Dist[Node][0] = 1;
    Depth[Node] = Dep;
    for(int i=0; i<Line[Node].size(); i++)
        if(Line[Node][i] != Par) setParent(Node, Line[Node][i], Dep+1);
}
 
int LCADist(int A, int B) {
    int Dis = 0;
    if(Depth[A] != Depth[B]) {
        if(Depth[A] < Depth[B]) swap(A, B);
        int Diff = Depth[A] - Depth[B];
        for(int i=0; Diff>0; i++) {
            if(Diff%2) {
                Dis += Dist[A][i];
                A = Parent[A][i];
            }
            Diff >>= 1;
        }
    }
    if(A != B) {
        for(int i=H; i>=0; i--)
            if(Parent[A][i] != 0 && Parent[A][i] != Parent[B][i]) {
                Dis += (Dist[A][i] + Dist[B][i]);
                A = Parent[A][i];
                B = Parent[B][i];
            }
        Dis += (Dist[A][0] + Dist[B][0]);
        A = Parent[A][0];
    }
    return Dis;
}
 
int main() {
    int N, M, A, B, C, temp, Sum = 0;
    scanf("%d", &N);
    temp = N;
    while(temp > 1) {
        temp /= 2;
        H++;
    }
    for(int i=0; i<N-1; i++) {
        scanf("%d %d", &A, &B);
        Line[A].push_back(B);
        Line[B].push_back(A);
    }
    setParent(0, 1, 0);
    for(int i=1; i<=H; i++)
        for(int j=2; j<=N; j++)
            if(Parent[j][i-1]) {
                Parent[j][i] = Parent[Parent[j][i-1]][i-1];
                Dist[j][i] = Dist[j][i-1] + Dist[Parent[j][i-1]][i-1];
            }
    scanf("%d", &M);
    scanf("%d", &A);
    for(int i=0; i<M-1; i++) {
        scanf("%d", &B);
        Sum += LCADist(A, B);
        A = B;
    }
    printf("%d", Sum);
}
 

 

풀이는 역시 위의 문제와 거의 동일한데, 초기 DFS 과정에서의 인접한 노드의 거리를 1로 설정해주고, 2^k 거리에 있는 도시 사이의 거리 갱신 부분은 그대로 두면 됩니다.

그리고 입력 부분 처리만 신경써주면 되는데, A를 먼저 하나 입력받고 이후의 M-1개의 B를 입력받아 LCADist(A, B)를 더해 인접한 두 입력의 거리를 구해주고, 이후 A를 B로 갱신해주는 과정만 반복적으로 적용해주면 됩니다.

 

 

 

m의 최댓값이 5000이기 때문에 시간이 길어지지 않을까 걱정하긴 했는데, 막상 풀이를 제출해보면 생각보다 시간이 거의 소요되지 않음을 알 수 있습니다.

만약 이 풀이 역시 TLE가 발생했다면, DP까지 활용해야 하기 때문에 문제의 난이도가 Diamond 이상으로 올라갔을 것 같습니다.

 

 

(+ 추가) 7742번 : Railway

 

이 문제도 역시 노드 사이의 거리를 구하는 문제인데, 난이도가 한 티어 높게 배정되어 있습니다.

 

#include <cstdio>
#include <vector>
#define MAX 100001
using namespace std;

int Parent[MAX][20], Dist[MAX][20] = {}, Depth[MAX], H = 0;
vector<pair<int, int>> Line[MAX];

void setParent(int Par, int Node, int Dis, int Dep) {
    if(!Line[Node].size()) return;
    Parent[Node][0] = Par;
    Dist[Node][0] = Dis;
    Depth[Node] = Dep;
    for(int i=0; i<Line[Node].size(); i++)
        if(Line[Node][i].first != Par) setParent(Node, Line[Node][i].first, Line[Node][i].second, Dep+1);
}

int findLCADis(int A, int B) {
    int Dis = 0;
    if(Depth[A] != Depth[B]) {
        if(Depth[A] < Depth[B]) swap(A, B);
        int Diff = Depth[A] - Depth[B];
        for(int i=0; Diff>0; i++) {
            if(Diff%2) {
                Dis += Dist[A][i];
                A = Parent[A][i];
            }
            Diff >>= 1;
        }
    }
    if(A != B) {
        for(int i=H; i>=0; i--)
            if(Parent[A][i] != 0 && Parent[A][i] != Parent[B][i]) {
                Dis += (Dist[A][i] + Dist[B][i]);
                A = Parent[A][i];
                B = Parent[B][i];
            }
        Dis += (Dist[A][0] + Dist[B][0]);
        A = Parent[A][0];
    }
    return Dis;
}

int main() {
    int N, M, A, B, C, temp;
    scanf("%d %d", &N, &M);
    temp = N;
    while(temp > 1) {
        temp /= 2;
        H++;
    }
    for(int i=0; i<N-1; i++) {
        scanf("%d %d %d", &A, &B, &C);
        Line[A].push_back({B, C});
        Line[B].push_back({A, C});
    }
    setParent(0, 1, 0, 0);
    for(int i=1; i<=H; i++)
        for(int j=2; j<=N; j++)
            if(Parent[j][i-1]) {
                Parent[j][i] = Parent[Parent[j][i-1]][i-1];
                Dist[j][i] = Dist[j][i-1] + Dist[Parent[j][i-1]][i-1];
            }
    for(int i=0; i<M; i++) {
        scanf("%d %d", &A, &B);
        printf("%d\n", findLCADis(A, B));
    }
}

 

입력부만 수정해서 같은 코드를 제출해도 정답을 인정받을 수 있습니다.