AT_arc151_b A < AP

Tag: 并查集，数学

题目描述

link

给定一个 $1\sim N$ 的排列 $P$，找到符合以下条件的 $A$ 数组的数量 $mod998244353$。

• 对于 $1\sim N$ 的每一个 $i$，$1\le A_i\le M$。
• $A$ 数组字典序小于 $(A_{P_1},A_{P_2},\cdots ,A_{P_n})$ 数组。

制約

• $1\le N\le 2\times {10}^5$
• $1\le M\le {10}^9$
• $1\le P_i\le N$
• $i\ne j⟹P_i\ne P_j$
• 入力はすべて整数

思路

我们令序列 $B=A_{p_1},A_{p_2},\cdots ,A_{p_n}$，题目要求 $A<B$，那可以枚举 $A$ 中第一个比 $B$ 小的位置。那么问题就可以转化成一个求连通块的问题。用并查集维护连通块。

当前有 $p$ 个连通块就代表有 $p$ 个可以填数的位置，
用快速幂就可以求出方案数。由于 $A<B$，所以 $(A_i,A_{p_i})$ 就有 ${\displaystyle \frac{m(m-1)}{2}}$ 种方案。

#include<bits/stdc++.h>
#define int long long
#pragma GCC optmize(3)
using namespace std;
const int mod=998244353;
int n,m,ans;
int cnt;
int p[200005],fa[200005];
int find(int x)
{
	if(fa[x]==x) return x;
	return fa[x]=find(fa[x]);
}
int pow(int a,int b,int p)
{
	int cnt=1;
	while(b)
	{
		if(b&1) cnt=cnt*a%p;
		a=a*a%p;
		b>>=1;
	}
	return cnt%p;
}
signed main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	cin>>n>>m;
	for(int i=1;i<=n;i++)
	{
		cin>>p[i];
		fa[i]=i;
	}
	cnt=n;
	for(int i=1;i<=n;i++)
	{
		int x=find(i),y=find(p[i]);
		if(x!=y)
		{
			ans=(ans+(pow(m,cnt-2,mod)*(m*(m-1)/2%mod))%mod)%mod;
			cnt--;
			fa[x]=y;
		}
	}
	cout<<ans<<"\n";
	return 0;
}