【牛客CSP-S提高组赛前集训营4】A - 复读数组【数论，数学】

题目大意:

题目链接：https://ac.nowcoder.com/acm/contest/1103/A

题目描述
有一个长为$n\times k$的数组，它是由长为n的数组$a_1,a_2,...,a_n$重复$k$次得到的。
定义这个数组的一个区间的权值为它里面不同的数的个数，现在，你需要求出对于这个数组的每个非空区间的权值之和。
答案对$10^9+7$取模。

---

思路：

考虑每一个数字会对多少个区间做贡献，容易发现，答案就是每一个数字能做贡献的区间数量之和。
又因为一个数字能做贡献的区间数量$=$总区间$-$没有该数字贡献的区间。
设$pos[i][j]$表示数字$i$第$j$次出现的位置。那么数字$i$的不做贡献的区间数量就是
$$count(1,pos[i][1]-1)+count(pos[i][m]+1,nk)+\sum _{i=2}^{m}count(pos[i][j-1]+1,pos[i][j]-1)$$
其中$count(x,y)=\frac{(r-l-1)(r-l-2)}{2}$
由于每重复的一段中的$count$之和,是一样的，所以我们对于${\sum }_{i=2}^{m}count(pos[i][j-1]+1,pos[i][j]-1)$只需要求出一次重复的情况，然后乘$k$即可。
注意还需要算上两个重复块之间的答案。如果一个块有$p$个数字$i$，那么答案就是$count(pos[i][p]+1,pos[i][1]+n-1)$。这种区间是有$k-1$个的。
时间复杂度$O(n\log n)$

---

代码：

#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;

const int N=100010,MOD=1e9+7;
int tot,a[N],b[N];
ll k,ans,cnt,n;
vector<int> pos[N];

ll power(ll x,ll k)
{
	ll ans=1;
	for (;k;k>>=1,x=x*x%MOD)
		if (k&1) ans=ans*x%MOD;
	return ans;
}

ll count(ll l,ll r)
{
	return (r-l+1)*(r-l+2)%MOD*power(2LL,MOD-2)%MOD;
}

int main()
{
	scanf("%lld%lld",&n,&k);
	for (int i=1;i<=n;i++)
	{
		scanf("%d",&a[i]);
		b[i]=a[i];
	}
	sort(b+1,b+1+n);
	tot=unique(b+1,b+1+n)-b-1;
	for (int i=1;i<=n;i++)
		pos[lower_bound(b+1,b+1+tot,a[i])-b].push_back(i);
	for (int i=1;i<=tot;i++)
	{
		int len=pos[i].size();
		cnt=(count(1,pos[i][0]-1)+count(pos[i][len-1]+1,n))%MOD;
		for (int j=1;j<len;j++)
			cnt=(cnt+count(pos[i][j-1]+1,pos[i][j]-1)*k%MOD)%MOD;
		cnt=(cnt+count(pos[i][len-1]+1,pos[i][0]+n-1)*(k-1)%MOD)%MOD;
		ans=(ans+count(1,(n*k)%MOD)-cnt)%MOD;
	}
	printf("%lld",(ans%MOD+MOD)%MOD);
	return 0;
}