C. Complex Market Analysis---双指针

题目链接

#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
#include <set>
#include <queue>
using namespace std;
#define int long long 
const int N = 1000010,mod = 1e9+7;
int a[N];
bool st[N];
int primes[N],pcnt=0;
set<int>S;
void prime()
{
	for(int i=2;i<N;i++)
	{
		if(!st[i])primes[pcnt++]=i;
		for(int j=0;primes[j]<=N/i;j++)
		{
			st[primes[j]*i]=true;
			if(i%primes[j]==0)break;
		}
	}
}

void solve()
{
	int n,e;cin>>n>>e;
	for(int i=0;i<n;i++)cin>>a[i],a[i]=a[i]==1?0:S.count(a[i])?1:2;
    int res=0;
    for(int c=0;c<e;c++)
    {
        vector<int>vec;
        for(int i=c;i<n;i+=e)vec.push_back(a[i]);
        int sum=0;
        for(int i=0,j=0;i<vec.size();i++)
        {
            sum+=vec[i];
            while(sum>1)
            {
                sum-=vec[j];
                j++;
            }
            res+=max(0ll,i-j);
        }
        
        sum=0;
        for(int i=0,j=0;i<vec.size();i++)
        {
            sum+=vec[i];
            while(sum>0)
            {
                sum-=vec[j];
                j++;
            }
            res-=max(0ll,i-j);
        }
    }
    cout<<res<<'\n';
}
 
signed main()
{	
	prime();
	for(int i=0;i<pcnt;i++)S.insert(primes[i]);
	int T;cin>>T;
	while(T--)solve();
}