[백준]개미(14942)

문제 링크 : [백준]개미(14942)

문제 설명

트리가 주어지고, 각 노드간의 거리가 주어집니다.
그때 각 노드별로 루트 노드를 향해 이동할 수 있는 비용이 주어집니다.
각 노드별로 루트에 가장 가깝게 갈 수 있는 노드의 번호를 출력하세요.

입력

N(노드의 갯수, 1 <= N <= 100,000)
E_i(노드별로 이동 할 수 있는 거리, 1 <= E_i <= 100,000)
Da_i Db_i Dc_i(노드 간 연결 정보, Da_i, Db_i 두 노드간 거리가 Dc_i 입니다. 1 <= Dc_i <= 10,000)

4
10
8
22
18
1 2 10
2 3 10
2 4 10

출력

1
2
1
2

카테고리

#누적합 #DFS #이분탐색

시간 복잡도 상한

O(NlogN)

해결 방법

1. DFS로 1번 노드(Root)에서 각 노드를 방문하며, 각 노드가 가진 비용만큼 위로 올라가면서 자신이 가진 비용보다 커지는 위치를 찾으며 답을 구하면 됩니다.
2. 그러나 최악의 경우 O(N²)이 되어 TLE가 됩니다.
3. 특정 노드에서 루트까지의 거리를 빠르게 구할 수 있다면 풀 수 있을 것 같습니다.
4. 각 노드를 기준으로, 자신의 부모 노드들까지의 거리를 알 수 있다면, upper_bound를 이용하여, O(logN)만에 각 노드에서 갈 수 있는 노드를 찾을 수 있습니다. 최종적으로 O(NlogN)에 해결 할 수 있습니다.
5. 한 노드에서 모든 부모노드까지의 거리를 구하려면 N²이 걸릴 것입니다.
6. DFS를 이용하면 단말 노드까지 오면서 자신의 부모에 대한 정보들을 유지하면서 올 수 있다는 점을 이용합니다.
7. 단말노드까지 오면서 거쳐온 정보들을 유지하면서, 단말노드에 도착하면, 단말노드를 기준으로 루트노드까지의 거리의 누적합들을 구합니다.
8. DFS에서 거꾸로 올라가면서, 단말노드에서 자신까지 오는 거리 + 자신이 갈 수 있는 거리를 가지고 upper_bound를 이용하여 자신이 갈 수 있는 부모 노드를 구합니다.

시간복잡도

O(NlogN)

코드

#include<iostream>
#include<algorithm>
#include<vector>
#include<deque>
#define M 100001
using namespace std;
typedef pair<int,int> P;
int n,d[M],a[M],i,x,y,z;
vector<P>v[M],w;
deque<int>d1,d2;
void r(int x,int u){
    int t=1;
    for(P k:v[x]){
        if(k.second==u)continue;
        t=0;
        w.push_back({k.first,x});
        r(k.second,x);
        w.pop_back();
    }
    if(t){
        d1.clear();d2.clear();
        d1.push_back(0);
        d2.push_back(x);
        for(int i=w.size()-1;i>=0;i--){
            d1.push_back(d1.back()+w[i].first);
            d2.push_back(w[i].second);
        }
        a[x]=d2[upper_bound(d1.begin(),d1.end(),d[x])-d1.begin()-1];
        d1.pop_front();d2.pop_front();
    }else{
        int c=d[x]+d1.front();
        d1.pop_front();d2.pop_front();
        if(d1[0]>c)a[x]=x;
        else a[x]=d2[upper_bound(d1.begin(),d1.end(),c)-d1.begin()-1];
    }
}
int main(){
    ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
    cin>>n;
    for(i=1;i<=n;i++)cin>>d[i];
    for(i=1;i<n;i++){
        cin>>x>>y>>z;
        v[x].push_back({z,y});
        v[y].push_back({z,x});
    }
    r(1,0);
    for(i=1;i<=n;i++)cout<<a[i]<<"\n";
}