M - Managing Difficulties

Problem

给定数列，求有多少个aj-ai=ak-aj，i<j<k。

Solution

枚举j和i，倍增差k，如果有，倍增查数量。时间复杂度n^2logn

Code

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>
using namespace std;
int t,n,a[2005],m,cnt[2005],all[2005][2005],b[2005];
struct ele{
    int x,id;
}e[2005];
bool cmp(const ele &x,const ele &y)
{
    if(x.x!=y.x)
    return x.x<y.x;
    return x.id<y.id;
}
int bzsearch(int mx,int pos)
{
    int nw=0;
    for(int i=mx;i>=0;i--)
    {
        if(nw+(1<<i)<=m&&b[nw+(1<<i)]<=pos)
            nw+=(1<<i);
    }
    return b[nw]==pos?nw:0;
}
int bzsearch1(int k,int pos)
{
    int mx=0,nw=cnt[k];
    while((1<<mx)<=cnt[k])
        mx++;
    for(int i=mx;i>=0;i--)
    {
        int nww=nw-(1<<i);
        if(nww>0&&all[k][nww]>pos)
            nw=nww;
    }
    return all[k][nw]<=pos?0:cnt[k]-nw+1;
}
void init()
{
    m=0;
    memset(cnt,0,sizeof(cnt));
}
int main()
{
    cin>>t;
    while(t--)
    {
        scanf("%d",&n);
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&a[i]);
            e[i].x=a[i];
            e[i].id=i;
        }
        sort(e+1,e+1+n,cmp);
        for(int i=1;i<=n;i++)
        {
            if(e[i].x!=e[i-1].x)
                m++;
            b[m]=e[i].x;
            all[m][++cnt[m]]=e[i].id;
        }
        long long ans=0;
        int mx=0;
        while((1<<mx)<=m) mx++;
        mx--;
        for(int i=2;i<n;i++)
        {
            for(int j=i-1;j>=1;j--)
            {
                int pos=a[i]+a[i]-a[j];
                int num=bzsearch(mx,pos);
                if(num)
                    ans+=bzsearch1(num,i);
            }
        }
        cout<<ans<<endl;
        init();
    }

    return 0;
}