题解 P6258 【[ICPC2019 WF]Hobson's Trains】

· · 题解

容易发现，原图是一个基环内向树森林。

对于每一棵基环树，首先把环去掉，考虑树。每一个树节点可以对自己不超过 k 级的祖先造成贡献，可以使用树上差分解决。

然后处理每个点对环的贡献。求出从每个点出发最先到达的环上的点 id_i，以及它离环的距离 len_i。那么该点对环造成的贡献是一个从点 id_i 开始，长 k-len_i+1 的连续段。仍然可以使用差分解决。

由于要处理的是一个环，会出现 3 \to 4 \to 5 \to 1 \to 2 这类终点编号小于起点编号的情况，所以要加上取模处理。

代码写得比较烂，各位凑合着看看吧。

# include <bits/stdc++.h>
# define rr
const int N=500010,INF=0x3f3f3f3f;
struct Edge{
    int to,next;
}edge[N];
int head[N],sum;
int d[N],loop[N],len[N],v[N],id[N];
// d[i] 含义见题面
//loop[i] 点所在环编号
//v[i] 点在环上的编号
int n,k;// 如题意
bool vis[N],ins[N]; // 是否访问过 是否在栈中
int sta[N],size;// 栈
std::vector <int> L[N]; // 环
std::vector <int> diff[N];// 环上差分
int cnt;// 环数量
int ans[N];// 每个店的最终答案
inline int read(void){
    int res,f=1;
    char c;
    while((c=getchar())<'0'||c>'9')
        if(c=='-')f=-1;
    res=c-48;
    while((c=getchar())>='0'&&c<='9')
        res=res*10+c-48;
    return res*f;
}
inline void add(int x,int y){
    edge[++sum].to=y;
    edge[sum].next=head[x];
    head[x]=sum;
    return;
}
void dfs(int i){
    vis[i]=true,ins[i]=true;
    sta[++size]=i;
    if(!vis[d[i]]){
        dfs(d[i]);
    }else if(ins[d[i]]){
        int now=size;
        ++cnt;
        while(sta[now]!=d[i]){
            L[cnt].push_back(sta[now]);
            --now; 
        }
        L[cnt].push_back(d[i]);
        std::reverse(L[cnt].begin(),L[cnt].end());
        for(rr int j=0;j<(int)L[cnt].size();++j){
            loop[L[cnt][j]]=cnt;
            v[L[cnt][j]]=j+1;
        }
        diff[cnt].resize((int)L[cnt].size()+2);// 处理环
    }
    ins[i]=false;
    --size;
    return;
}
void Tree_Diff(int i,int x){ // 树上差分
    sta[++size]=i;
    len[i]=len[sta[size-1]]+1;
    id[i]=x;
    ++ans[i];
    if((size-1)>k){
        --ans[sta[size-k-1]];
    }
    for(rr int j=head[i];j;j=edge[j].next){
        Tree_Diff(edge[j].to,x);
        ans[i]+=ans[edge[j].to];
    }
    --size;
    return;
}
int main(void){
    n=read(),k=read();
    for(rr int i=1;i<=n;++i){
        add(d[i]=read(),i);
    }
    for(rr int i=1;i<=n;++i){
        if(!vis[i]){
            dfs(i);
        }
    }
    for(rr int i=1;i<=n;++i){
        if(!loop[i]&&loop[d[i]]){
            size=0;
            len[0]=sta[0]=0;
            Tree_Diff(i,d[i]);
        }else if(loop[i]){
            id[i]=i;
        }
    }
    for(rr int i=1;i<=n;++i){
        if(len[i]<=k){
            int l=k-len[i];
            if(l>=(int)L[loop[id[i]]].size()-1){
                ++diff[loop[id[i]]][1];
            }else{
                int End=(v[id[i]]+l-1)%(int)L[loop[id[i]]].size()+1;
                if(v[id[i]]<=End){
                    ++diff[loop[id[i]]][v[id[i]]],--diff[loop[id[i]]][End+1];
                }else{
                    ++diff[loop[id[i]]][v[id[i]]],++diff[loop[id[i]]][1],--diff[loop[id[i]]][End+1];
                }
            }
        }
    }
    for(rr int i=1;i<=cnt;++i){
        int now=0;
        for(rr int j=0;j<(int)L[i].size();++j){
            now+=diff[i][j+1];
            ans[L[i][j]]+=now;
        }
    }
    for(rr int i=1;i<=n;++i){
        printf("%d\n",ans[i]);
    }
    return 0;
}