AT_arc151_b A < AP

Tag: 并查集，数学

题目描述

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

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

制約

• $ 1\ \leq\ N\ \leq\ 2\ \times\ 10^5 $
• $ 1\ \leq\ M\ \leq\ 10^9 $
• $ 1\ \leq\ P_i\ \leq\ N $
• $ i\ \neq\ j\ \implies\ P_i\ \neq\ 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})\) 就有 \(\dfrac{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;
}