[BOJ] 2472 체인점

굉장한 학생 문제랑 다익스트라의 짬뽕이다. 덕분에 코드가 장난아니게 길다.

각 정점이 최대 5개의 간선을 가지니 다익스트라의 $O(ElgV)$ 시간복잡도가 $O(VlgV)$로 바뀐다. 그러면 주어진 세 정점에 대해서 다익스트라를 한번씩 돌려볼 수 있고, 모든 정점에 대해 주어진 세 정점으로부터 각 정점으로부터의 거리 $dist_i$가 나오게 된다. 거리의 범위는 크지만 대소관계만을 비교하면 되므로 등위를 정렬을 통해 다시 매겨주자. 그러면 굉장한 학생 문제처럼 모든 정점에 대해 체인점을 세울수 있음/없음을 판별할 수 있다. 그 이후 들어오는 쿼리는 $O(1)$에 처리 가능하다.

 수행 시간은 다익스트라, 정렬, RMQ 등이 다 섞여서 결과적으로 $O(NlgN)$이다. 개인적으로 문제 자체도 굉장한 학생 문제를 모르면 접근하기 힘들고, 코딩 자체도 꽤나 까다로웠다고 생각함. 푸는 동안 재밌긴 했는데 다시 풀라고 하면 못 풀것 같다. ㅠㅠ

#include <cstdio>
#include <algorithm>
#include <vector>
#include <queue>
#include <tuple>
#include <cmath>
 
using namespace std;
 
const int INF=1e9+7;
int n,m,t;
int a[3];
 
typedef vector<int> vi;
typedef pair<int,int> pii;
typedef tuple<int,int,int,int> tp;
 
vector<vector<pii> > adj;
 
vi dijkstra(int st)
{
    vi dist(n+1,INF);
    priority_queue<pair<int,int> > Q;
    Q.push(make_pair(0,st));
    dist[st]=0;
    while(!Q.empty())
    {
        int hdist=-Q.top().first;
        int here=Q.top().second;
        Q.pop();
        if(hdist>dist[here]) continue;
        for(int i=0;i<adj[here].size();i++)
        {
            int there=adj[here][i].first;
            int tdist=adj[here][i].second+hdist;
            if(tdist>dist[there]) continue;
            dist[there]=tdist;
            Q.push(make_pair(-tdist,there));
        }
    }
    return dist;
}
 
typedef struct SegmentTree
{
    int n;
    vi tree;
    SegmentTree(int dn) :n(dn)
    {
        int tree_height=(int)ceil(log2(n));
        int tree_size=(1<<(tree_height+1));
        tree.assign(tree_size,INF);
    }
    void update(int st, int en, int idx, int node, int diff)
    {
        if(idx<st || en<idx) return;
        else if(st==en && st==idx)
        {
            tree[node]=diff;
            return;
        }
        int mid=(st+en)>>1;
        update(st,mid,idx,node*2,diff);
        update(mid+1,en,idx,node*2+1,diff);
        tree[node]=min(tree[node*2],tree[node*2+1]);
    }
    void update(int idx, int diff)
    {
        update(0,n-1,idx,1,diff);
    }
    int query(int st, int en, int le, int ri, int node)
    {
        if(ri<st || en<le) return INF;
        if(le<=st && en<=ri) return tree[node];
        int mid=(st+en)>>1;
        return min(query(st,mid,le,ri,node*2),query(mid+1,en,le,ri,node*2+1));
    }
    int query(int le, int ri)
    {
        return query(0,n-1,le,ri,1);
    }
}RMQ;
 
bool cmp(pii a, pii b)
{
    if(a.second==b.second) return a.first<b.first;
        else return a.second<b.second;
}
 
int main()
{
    scanf("%d",&n);
    for(int i=0;i<3;i++) scanf("%d",&a[i]);
    scanf("%d",&m);
    adj.assign(n+1,vector<pii>());
    for(int i=0;i<m;i++)
    {
        int t1,t2,tw;
        scanf("%d %d %d",&t1,&t2,&tw);
        adj[t1].push_back(make_pair(t2,tw));
        adj[t2].push_back(make_pair(t1,tw));
    }
    vi dist[3];
    for(int i=0;i<3;i++) dist[i]=dijkstra(a[i]);
 
    vector<pii> st[3];
    vi ra[3]; //등위를 저장
    for(int i=0;i<3;i++)
    {
        ra[i].assign(n+1,0);
        for(int j=1;j<=n;j++)
            st[i].push_back(make_pair(j,dist[i][j]));
        sort(st[i].begin(),st[i].end(),cmp);
        for(int j=0;j<st[i].size();j++)
            ra[i][st[i][j].first]=j;
 
    }
 
    vector<tp> VT;
    for(int i=1;i<=n;i++)
        VT.push_back(tie(ra[0][i],ra[1][i],ra[2][i],i));
    sort(VT.begin(),VT.end());
 
    vi okay(n+1);
    RMQ Q(n+1);
    for(int i=0;i<n;i++)
    {
        int da=get<0>(VT[i]), db=get<1>(VT[i]), dc=get<2>(VT[i]);
        int idx=get<3>(VT[i]);
        Q.update(db,dc);
        int mi=Q.query(0,db);
        if(mi<dc) okay[idx]=-1;
        else okay[idx]=1;
    }
 
    int rd;
    scanf("%d",&t);
    while(t--)
    {
        scanf("%d",&rd);
        printf("%s\n",okay[rd]==1?"YES":"NO");
    }
    return 0;
}