글 읽기 - centroid decomposition.. 왜 시간초과 날까요

dongwon0427 6년 전 0

O(n log2 n)인것 같은데.. set을 쓴게 문제가 될까요

#include <bits/stdc++.h>
using namespace std;
int n,k;
struct _tuple {
    int idx,val;
};
vector<_tuple> A[200005];
vector<int> vt;
int hashvt[200005];
int subsize[200005];
int disfromcen[200005];
int dfsvt[200005];
int whichsubtree[200005];
int depth[200005];
int dfs(int idx) {
    if(hashvt[idx]==0) return 0;
    if(dfsvt[idx]==1) return 0;
    dfsvt[idx] = 1;
    int ret = 1;
    for(int i=0;i<A[idx].size();i++) {
        if(hashvt[A[idx][i].idx]==1) ret+=dfs(A[idx][i].idx);
    }
    subsize[idx]=ret;
    return ret;
}
void dfs2(int idx,int dep,int key,int onedep) {
    if(hashvt[idx]==0) return;
    if(dfsvt[idx]==1) return;
    dfsvt[idx] = 1;
    disfromcen[idx]=dep;
    depth[idx]=onedep;
    whichsubtree[idx]=key;
    for(int i=0;i<A[idx].size();i++) {
        if(hashvt[A[idx][i].idx]==1) dfs2(A[idx][i].idx,dep+A[idx][i].val,key,onedep+1);
    }
}
void find_centroid() {
    dfs(vt[0]);
    for(int i=0;i<vt.size();i++) {
        dfsvt[vt[i]]=0;
    }
    int cent = vt[0];
    while(1) {
        int _max = 0;
        int maxidx = 0;
        for(int i=0;i<A[cent].size();i++) {
            if(hashvt[A[cent][i].idx]==0) continue;
            if(_max < subsize[A[cent][i].idx]) {
                _max = subsize[A[cent][i].idx];
                maxidx=A[cent][i].idx;
            }
        }
        if(_max*2 > subsize[cent]) {
            int tmpsize=subsize[maxidx];
            subsize[maxidx]=subsize[cent];
            subsize[cent]-=tmpsize;
            cent=maxidx;
        } else break;
    }
    int num=0;
    dfsvt[cent]=1;
    for(int i=0;i<A[cent].size();i++) {
        if(hashvt[A[cent][i].idx]==0) continue;
        dfs2(A[cent][i].idx,A[cent][i].val,++num,1);
    }
    for(int i=0;i<vt.size();i++) {
        dfsvt[vt[i]]=0;
    }
    //printf("%d",cent);
}
vector<int> subsnode[200005];
map<int,int> otherdis;
int ans = 2147483647;
void CenDec() {
    //for(int i=0;i<vt.size();i++) printf("%d ",vt[i]);
    //printf("\n");
    //for(int i=0;i<100000000;i++) i++;
    if(vt.size()<=1) {
        vt.clear();
        return;
    }
    for(int i=0;i<vt.size();i++) {
        hashvt[vt[i]]=1;
    }
    find_centroid();
    /*if(vt.size()==8) {
    for(int i=0;i<vt.size();i++) {
        printf("%d ",whichsubtree[vt[i]]);
    }
    printf("\n");
    }*/
    for(int i=0;i<vt.size();i++) {
        subsnode[whichsubtree[vt[i]]].push_back(vt[i]);
    }
    /*for(int i=0;true;i++) {
        if(subsnode[i].size()==0) break;
        for(int j=0;j<subsnode[i].size();j++) {
            printf("%d:%d ",subsnode[i][j],disfromcen[subsnode[i][j]]);
        }
        printf("\n");
    }*/
    for(int i=0;true;i++) {
        //printf("%d,,,,\n",i);
        if(subsnode[i].size()==0) break;
        for(int j=0;j<subsnode[i].size();j++) {
            //printf("%d~\n",k-disfromcen[subsnode[i][j]]);
            if(otherdis.find(k-disfromcen[subsnode[i][j]])!=otherdis.end()) {
                ans = min(ans,depth[subsnode[i][j]] + otherdis[k-disfromcen[subsnode[i][j]]]);
                //printf("YES, %d\n",ans);
            }
        }
        for(int j=0;j<subsnode[i].size();j++) {
            //printf("~%d ",subsnode[i][j]);
            if(otherdis.find(k-disfromcen[subsnode[i][j]])==otherdis.end()) otherdis[disfromcen[subsnode[i][j]]]=depth[subsnode[i][j]];
            else otherdis[disfromcen[subsnode[i][j]]] = min(depth[subsnode[i][j]], otherdis[disfromcen[subsnode[i][j]]]);
        }
        //printf("\n");
    }

    for(int i=0;i<vt.size();i++) {
        subsize[vt[i]]=0;
        disfromcen[vt[i]]=0;
        hashvt[vt[i]]=0;
        whichsubtree[vt[i]]=0;
        depth[vt[i]]=0;
    }
    vt.clear();
    otherdis.clear();
    vector<int> subsnode2[20005];
    for(int i=0;1;i++) {
        if(subsnode[i].size()==0) break;
        subsnode2[i]=subsnode[i];
        subsnode[i].clear();
    }

    for(int i=1;1;i++) {
        if(subsnode2[i].size()==0) break;
        for(int j=0;j<subsnode2[i].size();j++) {
            vt.push_back(subsnode2[i][j]);
        }
        CenDec();
    }
}
int main() {
    scanf("%d%d",&n,&k);
    for(int i=0;i<n-1;i++) {
        int t1,t2,t3;
        scanf("%d%d%d",&t1,&t2,&t3);
        _tuple tmp; tmp.idx=t2; tmp.val=t3;
        A[t1].push_back(tmp);
        tmp.idx=t1;
        A[t2].push_back(tmp);
    }
    for(int i=0;i<n;i++) vt.push_back(i);
    CenDec();
    if(ans==2147483647) printf("-1");
    else printf("%d",ans);
    return 0;
}