codeforces1261D1 Wrong Answer on test 233 (Easy Version)

https://codeforces.com/problemset/problem/1261/D1

首先k=1直接是0

对于一段长度为len相同的数字，他们向右移动一位，有len-1个是肯定不变的，每个位置都可以选1-k

那么我们考虑剩下的所有都是不同的数字的序列，每个位置的数字要么在之后+1，要么移动后-1，要么不变，这3种情况的方案数分别为1,1,k-2

注意，k=2的时候，一个不变的位置的方案是0，但是如果没有不变的位置时，方案数还是1的

那么我们可以直接枚举有多少个地方+1，有多少个地方-1，然后根据组合的计数原理推出组合数。

特殊情况：当所有数字都是相同过的时候，答案直接是0，最后30s过了样例交结果WA了，难受

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

const int mod=998244353;

int n,m,k,ad,eq,de,st,cnt;
int a[maxl],b[maxl];
long long ans;
long long jc[maxl],inv[maxl],mi[maxl];
char s[maxl];

inline long long qp(long long a,long long b)
{
	long long ans=1,cnt=a;
	while(b)
	{
		if(b&1)
			ans=ans*cnt%mod;
		cnt=cnt*cnt%mod;
		b>>=1;
	}
	return ans;
}

inline long long c(int n,int r)
{
	if(r>n || r<0) return 0;
	if(r==0 || r==n) return 1;
	return jc[n]*inv[r]%mod*inv[n-r]%mod;
}

inline void prework()
{
	scanf("%d%d",&n,&k);
	jc[0]=1;
	for(int i=1;i<=n;i++)
		jc[i]=jc[i-1]*i%mod;
	inv[n]=qp(jc[n],mod-2);
	for(int i=n-1;i>=0;i--)
		inv[i]=inv[i+1]*(i+1)%mod;
		
	if(k==1)
		ad=0,de=0,eq=1;
	else if(k==2)
		eq=0,ad=1,de=1;
	else
		ad=1,de=1,eq=k-2;
	mi[0]=1;
	for(int i=1;i<=n;i++)
		mi[i]=mi[i-1]*eq%mod;
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]),a[i+n]=a[i];
	for(int i=n;i<=2*n-1;i++)
	if(a[i]!=a[i+1])
	{
		st=i-n+1;
		break;
	}
	cnt=0;ans=1;
	for(int i=st;i<=st+n-1;i++)
	if(a[i]!=a[i-1])
		cnt++;
	ans=qp(k,n-cnt);
}

inline void mainwork()
{
	if(k==1 || cnt==1)
	{
		ans=0;
		return;
	}
	long long sum=0,t;
	for(int i=1;i<=cnt;i++)
		for(int j=i-1;j>=0;j--)
		if(i+j<=cnt)
		{
			t=c(cnt,i)*c(cnt-i,j)%mod*mi[cnt-i-j]%mod;
			sum=(sum+t)%mod;
		}
	ans=ans*sum%mod;
}

inline void print()
{
	printf("%lld",ans);
}

int main()
{
	int t=1;
	//scanf("%d",&t);
	for(int i=1;i<=t;i++)
	{
		prework();
		mainwork();
		print();
	}
	return 0;
}