백준[14942] 개미

문제 설명

문제 링크

http://icpc.me/14942

14942번: 개미

풀이

LCA와 거의 유사한 느낌으로 풀었다.

각 노드별로 모두 위로 올라가며 갈 수 있는 최대를 구한다면 쿼리 하나당 O(N)이 들어 시간이 터질 것이다.

그렇기 때문에 sparse table을 이용하여 시간을 줄여야 한다.

sparse table에는 dp[i][j]: 노드 j가 2^i만큼 조상 노드로 올라갔을 때의 비용과 조상 노드를 저장해놓는다.

이제 이 sparse table을 이용하면 쿼리당 시간을 O(logN)으로 줄일 수 있다.

코드

#include <bits/stdc++.h>
using namespace std;
using pii = pair<int, int>;

const int N = 1e5 + 5;
int energy[N];
vector<pii> v[N];
pii dp[20][N];

void dfs(int prv, int cur) {
    for (auto [nxt, cost] : v[cur]) {
        if (nxt == prv) continue;
        dp[0][nxt] = {cur, cost};
        dfs(cur, nxt);
    }
}

pii up(int cur, int cost) {
    for (int i = 19; i >= 0; i--) {
        auto [par, need_cost] = dp[i][cur];
        if (cost >= need_cost) return {par, cost - need_cost};
    }
    return {cur, cost};
}

int query(int cur, int cost) {
    while (1) {
        auto [nxt, last_cost] = up(cur, cost);
        if (nxt == 1 || cur == nxt) return nxt;
        cur = nxt;
        cost = last_cost;
    }
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    int n; cin >> n;
    for (int i = 1; i <= n; i++) cin >> energy[i];
    for (int i = 0; i < n - 1; i++) {
        int a, b, c; cin >> a >> b >> c;
        v[a].push_back({b, c});
        v[b].push_back({a, c});
    }
    dfs(-1, 1);
    dp[0][1] = {1, 0};
    for (int i = 1; i < 20; i++) for (int j = 1; j <= n; j++) {
        int par = dp[i - 1][j].first;
        dp[i][j] = {dp[i - 1][par].first, dp[i - 1][par].second + dp[i - 1][j].second};
    }
    for (int i = 1; i <= n; i++) cout << query(i, energy[i]) << '\n';
}