Jerry Kim

42는 ‘The Hitchhiker’s Guide to the Galaxy’라는 SF 소설에서 삶, 우주, 그리고 모든 것에 대한 궁극적인 질문에 대한 해답이다.

7 minute read

개미집은 n개의 방으로 구성되어 있으며, 이 방들은 1번부터 n번까지 번호가 부여되어 있다. 1번 방은 지면에 직접 연결되어 있는 방으로, 모든 개미는 이 방을 통해 지면으로 올라가고자 한다. 각 방은 서로 굴을 통해 연결되어 있으며, 굴을 이동하는 데는 굴의 길이만큼의 에너지가 소모된다. 개미들은 겨울잠에서 깨어나 지면으로 올라가기 위해 에너지를 사용하지만, 에너지가 부족하여 중간에 멈출 수 있다. 이 문제에서는 각 개미가 가진 에너지를 바탕으로, 도달할 수 있는 가장 1번 방에 가까운 방의 번호를 구하는 것이 목표이다.

문제 : https://www.acmicpc.net/problem/14942
 

접근 방식

이 문제는 트리 구조에서 각 노드(방) 간의 거리를 효율적으로 계산하고, 주어진 에너지로 최대로 가까운 1번 방에 도달할 수 있는 노드를 찾는 문제이다. 트리는 사이클이 없고, 모든 노드 간의 경로가 유일하기 때문에, 각 노드에서 1번 노드로 가는 경로가 명확하다.

효율적인 해결을 위해 다음과 같은 접근 방식을 사용하였다:

  1. 트리 구성 및 거리 계산:
    • 주어진 방과 굴 정보를 바탕으로 트리를 구성한다.
    • 깊이 우선 탐색(DFS)을 통해 각 방에서 1번 방까지의 누적 거리를 계산한다.
  2. 이진 승격(Binary Lifting) 기법 적용:
    • 이진 승격을 사용하여 각 노드의 2^k 번째 조상을 미리 계산한다.
    • 이를 통해 특정 에너지 내에서 도달 가능한 가장 가까운 조상을 빠르게 찾을 수 있다.
  3. 각 개미의 도달 가능한 방 찾기:
    • 각 개미의 에너지에 따라, 해당 에너지 내에서 도달할 수 있는 가장 가까운 방을 이진 승격을 통해 탐색한다.

이진 승격 기법을 사용함으로써, 각 개미의 도달 가능한 방을 O(log N) 시간에 찾을 수 있으며, 전체적으로 O(N log N)의 시간 복잡도로 문제를 해결할 수 있다.

C++ 코드와 설명

아래는 최적화된 C++ 코드와 각 라인에 대한 설명이다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int MAX = 100005;
const int LOG = 17; // 2^17 > 1e5

int n;
ll E_val[MAX];           // 각 방에 있는 개미의 에너지
ll dist_val[MAX];       // 1번 방부터 각 방까지의 거리
int up[LOG][MAX];       // 이진 승격을 위한 테이블
vector<pair<int, ll>> adj[MAX]; // 트리의 인접 리스트

// DFS를 통해 각 방의 부모와 거리를 계산
void dfs(int u, int p){
    up[0][u] = p;
    for(auto &[v, w] : adj[u]){
        if(v != p){
            dist_val[v] = dist_val[u] + w;
            dfs(v, u);
        }
    }
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    
    cin >> n;
    for(int u=1; u<=n; u++) cin >> E_val[u];
    for(int i=0; i<n-1; i++){
        int a, b;
        ll c;
        cin >> a >> b >> c;
        adj[a].emplace_back(b, c);
        adj[b].emplace_back(a, c);
    }
    
    // 1번 방을 루트로 DFS 수행
    dist_val[1] = 0;
    dfs(1, -1);
    
    // 이진 승격 테이블 구축
    for(int k=1; k<LOG; k++){
        for(int u=1; u<=n; u++){
            if(up[k-1][u] != -1){
                up[k][u] = up[k-1][up[k-1][u]];
            }
            else{
                up[k][u] = -1;
            }
        }
    }
    
    // 특정 방 u에서 에너지 E를 사용해 도달할 수 있는 가장 가까운 방 찾기
    auto get_ancestor = [&](int u, ll E) -> int {
        ll target_dist = dist_val[u] - E;
        if(target_dist <= 0){
            return 1;
        }
        int current = u;
        for(int k=LOG-1; k>=0; k--){
            if(up[k][current] != -1 && dist_val[up[k][current]] >= target_dist){
                current = up[k][current];
            }
        }
        return current;
    };
    
    // 각 개미에 대해 결과 출력
    for(int u=1; u<=n; u++){
        int v = get_ancestor(u, E_val[u]);
        cout << v << "\n";
    }
}

코드 설명

  1. 입력 처리:
    • n과 각 방에 있는 개미의 에너지 값을 입력받는다.
    • n-1개의 굴 정보를 입력받아 트리를 인접 리스트로 구성한다.
  2. DFS를 통한 거리 및 부모 노드 계산:
    • 1번 방을 루트로 설정하고 DFS를 수행하여 각 방까지의 거리(dist_val)와 부모 노드(up[0][u])를 계산한다.
  3. 이진 승격 테이블 구축:
    • 각 방에 대해 2^k 번째 조상을 미리 계산하여 up 테이블을 구축한다. 이는 나중에 특정 에너지 내에서 빠르게 조상을 찾기 위함이다.
  4. 조상 찾기 함수 get_ancestor:
    • 특정 방 u에서 에너지 E를 사용하여 도달할 수 있는 방을 찾는다.
    • 목표 거리를 dist_val[u] - E로 설정하고, 이진 승격을 통해 해당 거리 이상을 유지하면서 가능한 가장 높은 조상을 찾는다.
  5. 결과 출력:
    • 각 방의 개미에 대해 get_ancestor 함수를 호출하여 도달 가능한 방을 찾고, 그 방의 번호를 출력한다.

C++ without library 코드와 설명

이전 코드에서는 head 배열의 초기화 시점이 잘못되어 간선이 제대로 추가되지 않는 문제가 있었습니다. 이를 수정하여 올바르게 동작하는 코드를 제공하겠습니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

#include <stdio.h>
#include <stdlib.h>

typedef long long ll;

#define MAX 100005
#define LOG 17

ll E_val[MAX];
ll dist_val[MAX];
int up_table[LOG][MAX];
int head[MAX], to_arr[MAX*2], cost_arr[MAX*2], next_arr[MAX*2];
int cnt = 0;

// 간선 추가 함수
void add_edge(int a, int b, ll c){
    to_arr[cnt] = b;
    cost_arr[cnt] = c;
    next_arr[cnt] = head[a];
    head[a] = cnt++;
}

void dfs(int u, int p){
    up_table[0][u] = p;
    for(int i = head[u]; i != -1; i = next_arr[i]){
        int v = to_arr[i];
        ll w = cost_arr[i];
        if(v != p){
            dist_val[v] = dist_val[u] + w;
            dfs(v, u);
        }
    }
}

int main(){
    int n;
    scanf("%d", &n);
    
    // 에너지 값 입력
    for(int u=1; u<=n; u++) scanf("%lld", &E_val[u]);
    
    // head 배열 초기화 (간선 추가 전에 초기화해야 함)
    for(int i=1; i<=n; i++) head[i] = -1;
    
    // 간선 입력 및 추가
    for(int i=0; i<n-1; i++){
        int a, b;
        ll c;
        scanf("%d %d %lld", &a, &b, &c);
        add_edge(a, b, c);
        add_edge(b, a, c);
    }
    
    // DFS 수행하여 거리 및 부모 노드 계산
    dist_val[1] = 0;
    dfs(1, -1);
    
    // 이진 승격 테이블 구축
    for(int k=1; k<LOG; k++){
        for(int u=1; u<=n; u++){
            if(up_table[k-1][u] != -1){
                up_table[k][u] = up_table[k-1][up_table[k-1][u]];
            }
            else{
                up_table[k][u] = -1;
            }
        }
    }
    
    // 각 개미에 대해 도달 가능한 방 찾기
    for(int u=1; u<=n; u++){
        ll E = E_val[u];
        ll target_dist = dist_val[u] - E;
        if(target_dist <= 0){
            printf("1\n");
            continue;
        }
        int current = u;
        for(int k=LOG-1; k>=0; k--){
            if(up_table[k][current] != -1 && dist_val[up_table[k][current]] >= target_dist){
                current = up_table[k][current];
            }
        }
        printf("%d\n", current);
    }
}

코드 설명

  1. 입력 처리 및 간선 추가:
    • scanf를 사용하여 입력을 빠르게 처리한다.
    • 중요: 간선을 추가하기 전에 head 배열을 -1로 초기화한다. 이는 간선이 올바르게 추가되도록 보장하기 위함이다.
    • add_edge 함수를 통해 인접 리스트를 배열로 구현한다.
  2. DFS를 통한 거리 및 부모 노드 계산:
    • DFS를 수행하여 각 방까지의 거리(dist_val)와 부모 노드(up_table[0][u])를 계산한다.
  3. 이진 승격 테이블 구축:
    • 표준 라이브러리를 사용하지 않고, 배열을 통해 이진 승격 테이블을 구축한다.
  4. 조상 찾기 및 결과 출력:
    • 각 방의 개미에 대해 에너지에 따라 도달 가능한 방을 찾아 printf로 출력한다.

수정 사항:

  • head 배열 초기화 시점 변경: 간선을 추가하기 전에 head[i]를 -1로 초기화하도록 수정하였다. 이는 간선 추가 시 인접 리스트가 올바르게 구성되도록 보장한다.

Python 코드와 설명

아래는 최적화된 Python 코드와 각 라인에 대한 설명이다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

import sys
sys.setrecursionlimit(1 << 25)
from sys import stdin
def input():
    return sys.stdin.readline()

n = int(input())
E_val = [0] + [int(input()) for _ in range(n)]
adj = [[] for _ in range(n+1)]
for _ in range(n-1):
    a, b, c = map(int, input().split())
    adj[a].append((b, c))
    adj[b].append((a, c))

LOG = 17
up = [[-1]*(n+1) for _ in range(LOG)]
dist_val = [0]*(n+1)

def dfs(u, p):
    up[0][u] = p
    for v, w in adj[u]:
        if v != p:
            dist_val[v] = dist_val[u] + w
            dfs(v, u)

dfs(1, -1)

for k in range(1, LOG):
    for u in range(1, n+1):
        if up[k-1][u] != -1:
            up[k][u] = up[k-1][up[k-1][u]]

def get_ancestor(u, E):
    target_dist = dist_val[u] - E
    if target_dist <= 0:
        return 1
    current = u
    for k in reversed(range(LOG)):
        if up[k][current] != -1 and dist_val[up[k][current]] >= target_dist:
            current = up[k][current]
    return current

for u in range(1, n+1):
    v = get_ancestor(u, E_val[u])
    print(v)

코드 설명

  1. 입력 처리:
    • sys.stdin.readline()을 사용하여 빠르게 입력을 처리한다.
    • sys.setrecursionlimit을 설정하여 깊은 재귀를 허용한다.
  2. 트리 구성 및 DFS 수행:
    • 인접 리스트를 구성하고, DFS를 통해 각 방까지의 거리와 부모 노드를 계산한다.
  3. 이진 승격 테이블 구축:
    • 이진 승격 테이블을 리스트로 구현하여 각 노드의 2^k 번째 조상을 미리 계산한다.
  4. 조상 찾기 함수 get_ancestor:
    • 특정 방에서 에너지 내에서 도달 가능한 가장 가까운 조상을 이진 승격을 통해 찾는다.
  5. 결과 출력:
    • 각 방의 개미에 대해 도달 가능한 방을 출력한다.

결론

이번 문제는 트리 구조에서 각 노드 간의 거리를 효율적으로 계산하고, 이를 바탕으로 각 개미가 도달할 수 있는 방을 빠르게 찾는 문제였다. 이진 승격(Binary Lifting) 기법을 활용하여, 각 개미의 도달 가능한 방을 O(log N) 시간에 찾을 수 있었으며, 전체적으로 O(N log N)의 시간 복잡도로 문제를 해결할 수 있었다.

문제를 풀면서 트리 구조와 이진 승격 기법에 대한 이해를 더욱 깊게 할 수 있었으며, 이러한 기법들이 실제 문제에서 어떻게 적용될 수 있는지를 체험할 수 있었다. 추가적으로, Python과 C++ 두 가지 언어로 구현함으로써, 각 언어의 장단점을 비교해볼 수 있는 좋은 기회가 되었다. 앞으로도 다양한 트리 관련 문제를 풀면서 이진 승격을 비롯한 다양한 최적화 기법들을 익혀 나가고자 한다.

Tags: , , , , , , ,

Categories:

Source File: 2024-09-19-BOJ-14942.md

Updated:

Comments

You May Also Enjoy

[Algorithm] C++/Python 백준 5670번 : 휴대폰 자판

8 minute read

휴대폰에서 영단어를 입력할 때 버튼을 누르는 횟수를 최소화하기 위해 자동완성 기능을 활용하는 문제이다. 이 문제는 Trie(트라이) 자료 구조 또는 재귀적인 접근 방식을 사용하여 효율적으로 해결할 수 있다.

[Algorithm] C++/Python 백준 3679번 : 단순 다각형

6 minute read

이번 글에서는 주어진 점들로 자기 교차 없이 단순 다각형을 구성하는 문제인 백준 3679번 “단순 다각형”을 소개하고, 이를 해결하기 위한 접근 방식과 구현 방법을 살펴보겠다. 이 문제는 기하학적 알고리즘과 정렬을 활용하여 해결할 수 있으며, 효율적인 구현이 요구된다.

[Algorithm] C++/Python 백준 3653번 : 영화 수집

3 minute read

상근이는 영화 DVD를 수집하는 열성적인 수집가이다. 그는 자신의 DVD 콜렉션을 탑처럼 쌓아 보관한다. 영화를 보고 싶을 때마다 DVD의 위치를 찾아서, 쌓여 있는 콜렉션이 무너지지 않도록 조심스럽게 해당 DVD를 꺼낸다. 영화를 다 본 후에는 그 DVD를 가장 위에 놓는다.

[Algorithm] C++/Python 백준 3176번 : 도로 네트워크

5 minute read

N개의 도시와 그 도시들을 연결하는 N-1개의 도로로 이루어진 도로 네트워크가 있다. 모든 도시는 유일한 경로로 연결되어 있으며, 각 도로의 길이는 입력으로 주어진다.