백준 14942번: 개미 (C++)

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

14942번: 개미

---

Solved By: LCA, Sparse Table

 

1번 방까지 가야 하고, 경로가 n - 1개 있으며 한 방에서 다른 방으로 갈 수 있는 경로는 존재하며 유일하다 했으므로 이는 root node가 1인 트리를 의미합니다.

 

그리고, 각 방에서 주어진 에너지를 가지고, 1번 방까지 도달해야 하니 어떤 노드로부터 시작해서 1번 방으로 가는 길에 에너지를 다 쓰거나 남은 상태로 도착하게 되는 방은 어디인지를 구하는 문제였습니다.

 

LCA를 사용하면 풀 수 있을까요? 맞습니다.

 

1과 주어진 노드 사이의 경로가 몇 개있는지는 level 차이를 통해 구하고, 이 차이를 가지고서 개미를 각 노드로 이동시킵니다. 이동시키는 과정에서 Sparse Table을 활용하여 효율적으로 시간을 줄입니다.

 

그리고, 이동하면서 에너지를 다 썼는지를 알기 위해 거리 가중치 값이 항상 현재 보유하고 있는 에너지보다 낮아야 합니다.

물론, 현재 보유하고 있는 에너지는 도로를 지날 때마다 도로 가중치만큼 소모한다는 사실을 인지해야 합니다.

 

즉, 노드를 이동시키는 조건은 1을 향하며 보유하고 있는 에너지의 값을 경로 값들을 빼면서 에너지가 남아있는지 확인하는 것입니다. 다 썼다면 그 자리에서 멈추어야 합니다.

 

#include <iostream>
#include <vector>
#define MAX 100001
#define LOG_MAX 17
#define ll long long
using namespace std;

int n;
int energy[MAX];
int parent[MAX][LOG_MAX];
int dist[MAX][LOG_MAX];
int level[MAX];
vector<pair<int, ll>> adj[MAX];

void dfs(int now, int par){
    for(int i = 0; i < adj[now].size(); i++){
        int next = adj[now][i].first;
        ll nextCost = adj[now][i].second;
        
        if(next == par) continue;
        parent[next][0] = now;
        dist[next][0] = nextCost;
        level[next] = level[now] + 1;
        
        dfs(next, now);
    }
}

int getNum(int a, int b){
    int diff = level[b] - level[a];
    int cost = energy[b];
    
    for(int i = LOG_MAX - 1; i >= 0; i--){
        if(diff >= 1 << i && dist[b][i] <= cost){
            diff -= 1 << i;
            cost -= dist[b][i];
            
            b = parent[b][i];
        }
    }
    
    return b;
}

int main(){
    cin >> n;
    for(int i = 1; i <= n; i++) cin >> energy[i];
    
    for(int i = 0; i < n - 1; i++){
        int a, b, w; cin >> a >> b >> w;
        adj[a].push_back({b, w});
        adj[b].push_back({a, w});
    }
    
    level[1] = 1;
    dfs(1, 0);
    
    for(int j = 1; j < LOG_MAX; j++){
        for(int i = 1; i <= n; i++){
            parent[i][j] = parent[parent[i][j - 1]][j - 1];
            dist[i][j] = dist[i][j - 1] + dist[parent[i][j - 1]][j - 1];
        }
    }
    
    for(int i = 1; i <= n; i++){
        cout << getNum(1, i) << '\n';
    }
}