Newcoder 84 D.k进制数（数论+map）

Description

对于$k$进制数$x$，定义$d(x)$为$x$的各数位的和的$k$进制表示，如果结果超过一位，则继续重复执行各数位求和操作，直至结果为$1$位。

比如说，在$7$进制下，$d(3504_7)=d((3+5+0+4{)}_7)=d(15_7)=d((1+5{)}_7)=d(6_7)=6$

定义$x$为幸运的，当且仅当$d(x)=b$；

现在给定$k$进制下的$n$个数位$a_1{a}_2...a_n$，问其中有多少子串组成的数字是幸运的。

Input

第一行三个整数$k,b,n$

第二行$n$个整数$a_i$

$(2\le k\le 10^9,0\le b<k,1\le n\le 10^5,0\le a_i<k)$

Output

输出一个整数表示幸运的子串数。

Sample Input

10 5 6
3 2 0 5 6 1

Sample Output

5

Solution

显然$d(x)=(x-1)\%(k-1)+1$，$d(x)=0$当且仅当$x$每位都为$0$，首先统计连续为$0$的段的个数$res$，如果$b=0$则答案即为$res$，如果$b=k-1$，为方便统计，直接令$d(x)=x\%(k-1)$，如此求出的满足$d(x)=0$的子串本质上是$d(x)=0,k-1$的子串数量和，该值减去$res$即为答案

现在考虑求满足$d(x)=b$的子串个数，令$s_i=(a_1+...+a_i)\%(k-1)$，那么以$i$为右端点的合法区间$[j,i]$需要满足$s_i-s_{j-1}\equiv b(modk-1)$，用$map$统计每种前缀和出现次数即可

Code

#include<cstdio>
#include<map>
using namespace std;
typedef long long ll;
const int maxn=100005;
map<int,int>M;
int k,b,n,a[maxn];
int main()
{
	scanf("%d%d%d",&k,&b,&n);
	for(int i=1;i<=n;i++)scanf("%d",&a[i]);
	ll ans=0;
	for(int i=1,pre=1;i<=n;i++)
	{
		if(a[i])pre=i+1;
		else ans+=i-pre+1;
	}
	if(b==0)
	{
		printf("%lld\n",ans);
		return 0;
	}
	if(b==k-1)ans=-ans,b=0;
	else ans=0;
	M[0]=1;
	for(int i=1,res=0;i<=n;i++)
	{
		res=(res+a[i])%(k-1);
		ans+=M[(res-b+k-1)%(k-1)];
		M[res]++;
	}
	printf("%lld\n",ans);
	return 0;
}