[BOJ 14942 // C++] 개미

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

 

이번에 볼 문제는 백준 14942번 문제인 개미이다.
문제는 아래 링크를 확인하자.

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

14942번: 개미

각 개미가 있는 위치에서 주어진 에너지로 갈 수 있는 루트에 가장 가까운 노드를 찾는 문제이다.

 

각 노드마다 2^k개의 노드만큼씩 루트방향으로 움직일 때 필요한 에너지를 sparse table의 형태로 전처리해둔다면 문제를 간단히 해결할 수 있다.

 

아래는 제출한 소스코드이다.

#include <iostream>
#include <vector>
using namespace std;

vector<pair<int, int>> G[100001];
int energy[100001];
int sparsetable[100001][20];
int tableparent[100001][20];

void dfs(int node, int parent, int dist) {
    int tparent = parent;
    for (int i = 0; i < 20; i++) { // 2^i 칸 움직일 때 드는 에너지(dist)
        sparsetable[node][i] = dist;
        tableparent[node][i] = tparent;
        dist += sparsetable[tparent][i];
        tparent = tableparent[tparent][i];
    }
    vector<pair<int, int>>::iterator iter = G[node].begin();
    for (; iter != G[node].end(); iter++) {
        int current = (*iter).first;
        int currentdist = (*iter).second;

        if (current != parent)
            dfs(current, node, currentdist);
    }
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);

    int N;
    cin >> N;

    for (int i = 1; i <= N; i++) {
        int x;
        cin >> x;
        energy[i] = x;
    }

    for (int i = 1; i < N; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        G[a].push_back({ b,c });
        G[b].push_back({ a,c });
    }

    for (int i = 0; i < 20; i++) {
        tableparent[1][i] = 1;
    }

    dfs(1, 1, 0);

    for (int i = 1; i <= N; i++) {
        int currentnode = i;
        int currentenergy = energy[i];

        while (currentnode != 1) {
            bool chk = 0;
            for (int j = 19; j >= 0; j--) {
                if (sparsetable[currentnode][j] <= currentenergy && sparsetable[currentnode] != 0) {
                    currentenergy -= sparsetable[currentnode][j];
                    currentnode = tableparent[currentnode][j];
                    chk = 1;
                    break;
                }
            }
            if (!chk)
                break;
        }
        cout << currentnode << '\n';
    }
    return 0;
}