2021牛客暑期多校训练营5 1009.Array（线段树/树状数组）

题目链接：https://acm.hdu.edu.cn/showproblem.php?pid=7020

分析

如果用线段树会超时，可以用树状数组，参考：https://www.luogu.com.cn/problem/solution/P4062

线段树代码（超时）：

//线段树维护等差数列
#include<bits/stdc++.h>
using namespace std;
const int N=1e6+10;
vector<int>g[N];
int a[N],ans[N];
struct node
{
    int l,r;
    int sum,a,d;
}tr[N<<3];
void pushup(int k)
{
    tr[k].sum=tr[k<<1].sum+tr[k<<1|1].sum;
    return ;
}
void pushdown(int k)
{
    if(tr[k].a==0&&tr[k].d==0) return ;
    int mid=(tr[k].l+tr[k].r)/2;
    tr[k<<1].a+=tr[k].a;
    tr[k<<1|1].a+=tr[k].a+(mid-tr[k].l+1)*tr[k].d;
    tr[k<<1].d+=tr[k].d;
    tr[k<<1|1].d+=tr[k].d;
    tr[k].a=tr[k].d=0;
    return ;
}
void build(int k,int l,int r)
{
    tr[k].l=l,tr[k].r=r;
    if(l==r)
    {
        tr[k].sum=tr[k].a=tr[k].d=0;
        return ;
    }
    int mid=(l+r)/2;
    build(k<<1,l,mid);
    build(k<<1|1,mid+1,r);
}
void update(int k,int dl,int dr,int a,int d)
{
    if(tr[k].l>=dl&&tr[k].r<=dr)
    {
        tr[k].a+=a+(tr[k].l-dl)*d;
        tr[k].d+=d;
        tr[k].sum+=(a+(tr[k].l-dl)*d+a+(tr[k].r-dl)*d)*(tr[k].r-tr[k].l+1)/2;
        return ;
    }
    pushdown(k);
    int mid=(tr[k].l+tr[k].r)/2;
    if(dl<=mid) update(k<<1,dl,dr,a,d);
    if(dr>mid) update(k<<1|1,dl,dr,a,d);
    pushup(k);
}
int query(int k,int dl,int dr)
{
    if(tr[k].l>=dl&&tr[k].r<=dr) return tr[k].sum;
    pushdown(k);
    int mid=(tr[k].l+tr[k].r)/2;
    int ans=0;
    if(dl<=mid) ans=query(k<<1,dl,dr);
    if(dr>mid) ans+=query(k<<1|1,dl,dr);
    return ans;
}
int main()
{
    int t;
    scanf("%d",&t);
    while(t--)
    {
        int n;
        scanf("%d",&n);
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&a[i]);
            g[a[i]].push_back(i);
        }
        build(1,1,N*2);
        int res=0;
        sort(a + 1, a + 1 + n);
        int m = unique(a + 1, a + 1 + n) - a - 1;
        for(int i=1;i<=m;i++)
        {
            //if(ans[a[i]]==1) continue;
            g[a[i]].push_back(n+1);
            int st=0;
            for(int j=0;j<g[a[i]].size();j++)
            {
                int now=g[a[i]][j]-1;
                int d=now-st;
                int x=j*2-now+N,y=j*2-st+N;
                res+=query(1,x-1,y-1);
                update(1,x,y,1,1);
                update(1,y+1,2*N,y-x+1,0);
                st=g[a[i]][j];
            }
            st=0;
            for(int j=0;j<g[a[i]].size();j++)
            {
                int now=g[a[i]][j]-1;
                int d=now-st;
                int x=j*2-now+N,y=j*2-st+N;
                update(1,x,y,-1,-1);
                update(1,y+1,2*N,-(y-x+1),0);
                st=g[a[i]][j];
            }
            //ans[a[i]]=1;
        }
        printf("%d\n",res);
        for(int i=1;i<=n;i++) 
        {
            g[a[i]].clear();
            ans[a[i]]=0; 
        }
    }
    system("pause");
    return 0;
}

---

用树状数组维护三阶前缀和，推出公式

那么可以开三个树状数组来维护$d_k$，$d_{k}k$，$d_k{k}^2$。

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int INF = 0x3f3f3f3f;
const LL mod = 1e9 + 7;
const int N = 1000005;
int ans[N];
LL c1[N * 4], c2[N * 4], c3[N * 4];
LL sum(int x) {
    LL res = 0;
    for (int i = x; i > 0; i -= i & -i) {
        res += c1[i] * (x + 2) * (x + 1) - c2[i] * (2 * x + 3) + c3[i];
    }
    return res / 2;
}
void add(int x, LL d, int n) {
    for (int i = x; i <= n; i += i & -i) {
        c1[i] += d;
        c2[i] += d * x;
        c3[i] += d * x * x;
    }
}
int a[N];
vector<int> b[N];
int main() {
    int t;
    scanf("%d",&t);
    while(t--)
    {
        int n;
        scanf("%d", &n);
        for (int i = 1; i <= n; i++) 
        {
            scanf("%d", &a[i]);
            b[a[i]].push_back(i);
        }
        LL res = 0;
        for (int i = 1; i <=n; i++)
        {
            if(ans[a[i]]==1) continue;
            b[a[i]].push_back(n + 1);
            int last = 0;
            for (int j = 0; j < b[a[i]].size(); j++) 
            {
                int y = 2 * j - last + N, x = 2 * j - (b[a[i]][j] - 1) + N;
                res += sum(y - 1) - (x >= 3 ? sum(x - 2) : 0);
                add(x, 1, 2 * N );
                add(y + 1, -1, 2 * N);
                last = b[a[i]][j];
            }
            last = 0;
            for (int j = 0; j < b[a[i]].size(); j++) {
                int y = 2 * j - last + N, x = 2 * j - (b[a[i]][j] - 1) + N;
                add(x, -1, 2 * N);
                add(y + 1, 1, 2 *N);
                last = b[a[i]][j];
            }
            ans[a[i]]=1;
        }
        printf("%lld\n", res);
        for(int i=1;i<=n;i++)
        {
            b[a[i]].clear();
            ans[a[i]]=0;
        }
    }
    //system("pause");
    return 0;
}