持续创作，加速成长！这是我参与「掘金日新计划 · 10 月更文挑战」的第11天，点击查看活动详情

解题记录

K 复合函数

题目大意

给定正整数 $n$，并记 $I_n = {1, 2, . . . , n}$。 给定一个定义在 $I_n$ 上的函数 $f$，满足对于任意 $x \in I_n$ 有 $f(x) \in I_n$。

对于任意正整数 $k$，当 $k > 1$ 时定义 $f^{k} (x) = f ( f ^{k−1} (x) )$ ，特殊地，当 $k = 1$ 时定义 $f ^1 (x) = f(x)$。

现有 $q$ 组询问，每组询问给定两个正整数 $a,b$，试求出有多少个 $x \in I_n$ 满足 $f ^a (x) = f ^b (x)$。

思路

考虑一张 $n$ 个点 $n$ 条边的有向图 $G$，其中点 $x$ 向 $f(x)$ 连边。

则样例 2 建出来的图如下：

容易发现 $G$ 由若干个基环树构成。 对于一组询问，基环树之间是独立的，接下来我们只考虑一 棵基环树对答案的贡献。

我们令 $a$ 取 $a,b$ 中的较小值，$b$ 取 $a,b$ 中的较大值，原问题可以转化为有多少 $x$ 满足 $f^a(x)=f^{b-a}(f^a(x))$。

记基环树的环长为 $c$，对于基环树内的点 $x$，$f ^a (x) = f ^b (x)$ 成立，当且仅当以下两个条件至少一个成立：

• $a = b$
• $c \mid (b-a)$ 且 $f ^a (x)$ 在环上

注意到条件中只需要关注环长以及基环树上各点的深度，因此我们需要统计出环长为 $c$ 的基环树中到环距离为 $d$ 的点的数量为 cnt[c][d]。对 cnt 数组做前缀和就能求出深度不超过 $d$ 的点的数量 sum[c][d]。可以发现，$c\in [1,n]$，但是 $G$ 中基环树的环长最多只有 $O(\sqrt n)$ 种。可以对环长进行离散化优化空间。

对于每组询问，
若 $a=b$，原式转化为$f ^a (x) = f ^a (x)$，显然答案是 $n$。
否则，答案是 $\sum_{c_i\mid (b-a)} sum[c_i][a]$。

时间复杂度 $O(n \sqrt n)$。

代码

#include <stdio.h>
#include <vector>
#include <algorithm>
#include <math.h>
#include <string.h>
#include <map>
using namespace std;
const int N=100001;
int f[N],n,q;
int bel[N],d[N],cnt[N],idx,vis[N],stk[N],tot,t;
int sum[501][N];
map<int,int> mp;
long long a,b;
int main()
{
	scanf("%d",&n);
	for (int i=1;i<=n;++i) scanf("%d",&f[i]);
	for (int i=1;i<=n;++i)
	{
		if (bel[i]) continue;
		for (int j=i;vis[j]==0&&bel[j]==0;j=f[j])
		{
			stk[++tot]=j;
			vis[j]=1;
		}
		t=f[stk[tot]];
		if (bel[t])
		{
			while (tot>=1)
			{
				bel[stk[tot]]=bel[t];
				d[stk[tot]]=d[f[stk[tot]]]+1;
				tot--;
			}
			continue;
		}
		++idx;
		while (tot>=1)
		{
			cnt[idx]++;
			bel[stk[tot]]=idx;
			if (stk[tot]==t)
			{
				tot--;
				break;
			}
			tot--;
		}
		while (tot>=1)
		{
			bel[stk[tot]]=bel[t];
			d[stk[tot]]=d[f[stk[tot]]]+1;
			tot--;
		}
	}
	tot=0;
	memset(vis,0,sizeof(vis));
	for (int i=1;i<=n;++i)
	{
		t=mp[cnt[bel[i]]];
		if (!t)
		{
			++tot;
			mp[cnt[bel[i]]]=t=tot;
			vis[tot]=cnt[bel[i]];
		}
		sum[t][d[i]]++;
	}
	for (int i=1;i<=tot;++i)
		for (int j=1;j<=n;++j) sum[i][j]+=sum[i][j-1];
	long long a,b,ans;
	for (scanf("%d",&q);q--;)
	{
		scanf("%lld%lld",&a,&b);
		if (a>b) swap(a,b);
		if (a==b) ans=n;
		else
		{
			ans=0;
			for (int i=1;i<=tot;++i)
			{
				if ((b-a)%vis[i]) continue;
				if (a>n) ans+=sum[i][n];
				else ans+=sum[i][a];
			}
		}
		printf("%lld\n",ans);
	}
	return 0; 
}