【基环树】HDU5329 Question for the Leader

$Souce:$2015 Multi-University Training Contest 4
$Problem:$一棵$n=1e5$的基环树，对于$k=1...n$，问是否能将基环树切割成k份大小相同的连通块。
$Idea:$如果只是一棵树切成k份，当且仅当$k$个点$sz[u] ≡ 0 mod (n/k) $时成立。那么只剩下处理环上的情况，需要最优的切法，使得断环为链后贡献最大，可以直接模拟下，维护下前缀和的数量线性递推。
$Code:$

#include<bits/stdc++.h>
using namespace std;

#define lc o<<1
#define rc o<<1|1
#define fi first
#define se second
#define pb emplace_back
#define ALL(X) (X).begin(), (X).end()
#define bcnt(X) __builtin_popcountll(X)
#define CLR(A, X) memset(A, X, sizeof(A))
#pragma comment(linker, "/STACK:1024000000,1024000000")

using DB = double;
using LL = long long;
using PII = pair<int, int>;
const int N = 2e5+10;
const int INF = 0x3f3f3f3f;

int sz[N], fa[N], vis[N], b[N], cnt[N], pre[N];
vector<int> G[N];

void dfs(int u, int fa) {
    sz[u] = 1;
    for(int &v:G[u]) if(vis[v]!=2 && v!=fa) {
        dfs(v, u);
        sz[u] += sz[v];
    }
}

int solve(int n, int d) {
    int sum = 0, k = 0, u;
    cnt[0] = 0;
    for(int i = 1; i <= n; i++) {
        cnt[i] = 0;
        if(vis[i] == 3) vis[i] = 2;
        if(vis[i] != 2) {
            if(sz[i]%d == 0) sum++;
        }
        else u = i;
    }
    while(vis[u] == 2) {
        k++; b[k] = u;
        pre[k] = (sz[u]%d+pre[k-1])%d;
        cnt[pre[k]]++;
        vis[u] = 3, u = fa[u];
    }

    if(cnt[0]+sum == n/d) return 1;
    int last = k;
    for(int i = 1; i < k; i++) {
        int u = b[i], tmp = pre[i];
        cnt[pre[i]]--;
        pre[i] = (sz[u]+pre[last])%d;
        cnt[pre[i]]++;
        last = i;
        if(cnt[tmp]+sum == n/d) return 1;
    }
    return 0;
}

void work() {
    int n;
    if(scanf("%d", &n) == -1) exit(0);
    for(int i = 1; i <= n; i++) {
        vis[i] = 0;
        G[i].clear();
    }
    for(int i = 1; i <= n; i++) {
        int fa; scanf("%d", &fa);
        G[fa].pb(i), G[i].pb(fa);
        ::fa[i] = fa;
    }
    int u = 1;
    while(vis[u] == 0) vis[u] = 1, u = fa[u];
    while(vis[u] == 1) vis[u] = 2, u = fa[u];
    for(int i = 1; i <= n; i++) if(vis[i] == 2) {
        dfs(i, 0);
    }
    int ans = 0;
    for(int i = 1; i <= n; i++) if(n%i == 0) {
        ans += solve(n, i);
    }
    printf("%d\n", ans);
}

int main() {
    for(;;) work();
    return 0;
}