kmjp's blog

競技プログラミング参加記です

TopCoder SRM 749 Div2 Hard CountSubarrays

まぁこれは典型問題だし900ptでもしょうがないよね。
https://community.topcoder.com/stat?c=problem_statement&pm=15301

問題

整数列Aと整数Xが与えられる。
素数P=744,444,499に対し、Aの部分列A[L...R]の積をPで割った余りがXとなるような部分列は何通りか。

解法

Xが0かどうかで場合分けする。

  • X=0の場合、部分列に1個でも0が含まれれば積は0で、1個も含まなければ0にはならない。
    • 「0が1個でも含まれる区間の数」を数えるのは面倒なので、部分列の全組み合わせから、0が1個も含まれない区間の数を引こう。
  • X!=0の場合、良くある「区間和がXとなる問題の区間の数え上げ」を、累積和の代わりに積と除算で行うことを考えよう。
    • B(x) = prod(A[0..x])とすると、Rを小さい方からずらしながらB(R)を求めていき、そのつどB(R)/B(L-1)=XとなるようなLを数え上げよう。これは過去の登場した累積積B(L-1)のうち、B(R)/Xとなるものを列挙すればよい。
    • ただし、A[L...R]の途中に0が入っているものは不可なのでA[R]=0となるRを処理した際のそれ以前の累積積は忘れるようにすること。
ll mo=744444499;
ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}
ll rev(ll a) {
	ll r=1, n=mo-2; a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

class CountSubarrays {
	public:
	long long count(vector <int> A, int n, int a, int b, int m, int X) {
		vector<ll> B;
		FORR(v,A) B.push_back(v);
		for(int i=A.size();i<n;i++) B.push_back((B[i-A.size()]*a+b)%m);
		ll ret=0;
		if(X==0) {
			int num=0;
			ret=1LL*n*(n+1)/2;
			FORR(b,B) {
				if(b==0) num=0;
				else num++;
				ret-=num;
			}
		}
		else {
			map<ll,int> mp;
			mp[1]=1;
			ll cur=1;
			FORR(b,B) {
				if(b==0) {
					mp.clear();
					mp[1]=1;
					cur=1;
				}
				else {
					cur=cur*b%mo;
					ret+=mp[cur*modpow(X)%mo];
					mp[cur]++;
				}
			}
		}
		return ret;
	}
}

まとめ

最初X=0のケースを見落としていて、「さすがにDiv2Hardでもこの難易度は無いでしょ…」と思った。