【树状数组维护三阶前缀和】洛谷 P4062 Yazid的新生舞会

题目链接

题链

题解

区间众数的个数 $>$ 区间长度一半 称这个区间有主元素，主元素就是这个众数；

题意：求数组中有多少个区间有主元素;

考虑一个子问题：每一种数作为主元素的贡献；

例如给定数组 $p=[3,2,1,3,3,2]$，并考虑 $3$ 作为主元素的贡献；

我们可以将是数字 $3$ 的记为 $1$ ，不是的记为 $-1$；

	$1$	$2$	$3$	$4$	$5$	$6$
$p$	$3$	$2$	$1$	$3$	$3$	$2$
$w$	$1$	$-1$	$-1$	$1$	$1$	$-1$

于是子问题变成求 $w$ 中有多少个区间和大于 $0$，例如区间 $[3,5]$ 的和为 $1$ ，所以这个区间可行，而区间 $[2,5]$ 的和为 $0$，所以这个区间不可行；

我们对 $w$ 求一个前缀和，记为 $d$，如果 $d_r-d_{l-1}>0$，也就是对于 $d_r$ 查找 $r$ 之前有多少个比 $d_r$ 小的数字，那么区间 $[l,r]$ 满足要求；

	$0$	$1$	$2$	$3$	$4$	$5$	$6$
$p$	无	$3$	$2$	$1$	$3$	$3$	$2$
$w$	$0$	$1$	$-1$	$-1$	$1$	$1$	$-1$
$d$	$0$	$1$	$0$	$-1$	$0$	$1$	$0$

例如对于 $r=5$ ，$r$ 前比 $d_r$ 小的数字有 $4$ 个，分别是 $l-1=[4,3,2,0]$ 的时候，也就是区间 $[5,5],[4,5],[3,5],[1,5]$ 满足要求；

于是可以对每一种数字，求出 $w$，求出 $d$，枚举 $d_i$，查找 $i$ 前有多少个比 $d_i$ 小的数字，用树状数组维护权值，显然这样的复杂度是 $O(n^{2}logn)$的，不可取；

我们记录下每一种数字在数组中的位置，例如数字 $3$，它的位置有 $po{s}_3=[1,4,5]$，如果能在 $O(po{s}_i.size())$ 的时空复杂度下(可以加个$log$)求出这种数字的贡献，那么枚举每一种数字算贡献整体就是 $O(nlogn)$ 的了；

$d$ 是连续的，$po{s}_3=[1,4,5]$ 可以将 $d$ 分为 $4$ 个部分，分别是 $[0,0],[1,3],[4,4],[5,6]$ ，这四个部分都是等差数列，对于$[1,3]$区间内的每个值，在该值前且比该值小的数字只能来自于$[1,3]$这部分之前，不能来自于 $[1,3]$ 这部分本身；

设法将每一部分同时处理，如有一部分 $P_i$，$P_i$这部分的答案来自于 $P_i$ 之前的部分，对 $P_i$ 这部分是一个等差数列 $[y,y-1,y-2,\dots ,x]$，那么 $P_i$ 部分插入到树状数组中就是区间 $[x,y]$ 加 $1$；

设 $B_i$ 为权值的前缀和，$B_i=\sum _{j=1}^i{C}_j$，那么对于 $P_i$ 部分内的每一个值 $d_i$，$d_i$ 的贡献就等于 $B_{d_i-1}$，那么对于这一整个部分 $P_i$，值为 $[x,y]$，整体贡献就是 $\sum _{i=x-1}^{y-1}{B}_i$，设 $A_i$ 为权值的前缀和的前缀和，也就是 $B_i$ 的前缀和 $\sum _{j=1}^i{B}_j$，那么整体贡献就是 $A_{y-1}-A_{x-2}$；

于是将问题转变为了，区间 $[x,y]$ 加 $1$，和求二阶前缀和；

区间 $[x,y]$加一个数，求一阶前缀和就是区间加数，区间求和的洛谷线段树模板题1，当然可以用树状数组实现；

和模板题同样将权值 $C_i$差分，差分后的数组位 $D_i$，于是问题转变为单点修改，求三阶前缀和；

$A_x=\sum _{i=1}^x\sum _{j=1}^i\sum _{k=1}^j{D}_k$
$A_x$
$=B_1+B_2+B_3+\cdots +B_x$
$=C_1+(C_1+C_2)+(C_1+C_2+C_3)+\cdots +(C_1+C_2+C_3+\cdots +C_x)$
$=x{C}_1+(x-1)C_2+(x-2)C_3+\cdots +(x-n+1)C_n+\cdots +C_x$
$=x{D}_1+(x-1)(D_1+D_2)+(x-2)(D_1+D_2+D_3)+\cdots +(x-n+1)(D_1+D_2+\cdots +D_n)+\cdots +(D_1+D_2+D_3+\cdots +D_n)$
$=\frac{(x+1)x}{2}{D}_1+\frac{x(x-1)}{2}{D}_2+\frac{(x-1)(x-2)}{2}{D}_3+\cdots +\frac{(x-n+2)(x-n+1)}{2}{D}_n+\cdots +D_x$
$=\sum _{n=1}^x\frac{(x-n+2)(x-n+1)}{2}{D}_n$
$=\sum _{n=1}^x{n}^2{D}_n+\frac{-2x+3}{2}\sum _{n=1}^{x}n{D}_n+\frac{x^2+3x+2}{2}\sum _{n=1}^x{D}_n$

用三个树状数组维护 $i^2{D}_i$，$i{D}_i$，$D_i$即可；

#include <bits/stdc++.h>
using namespace std;
#define LL long long
#define MS 2000009

LL n,m;
LL a[MS];
vector<LL > vc[MS];
LL b[MS], tot;
LL p[4][MS];
LL ans;

LL lowbit(LL x){
    return x&(-x);
}

void add(LL pos, LL val){
    for(LL x=pos;pos<MS;pos+=lowbit(pos)){
        p[1][pos] += x*x*val;
        p[2][pos] += x*val;
        p[3][pos] += val;
    } 
}

LL query(LL pos){
    LL ans = 0;
    for(LL x=pos;pos;pos-=lowbit(pos)){
        ans += p[1][pos];
        ans += (-2*x-3)*p[2][pos];
        ans += (x*x+3*x+2)*p[3][pos];
    } 
    return ans/2;
}

void clear(){
    for(LL i=1;i<=tot;i++) vc[b[i]].clear();
    tot = 0;
    ans = 0;
}

void solve(){
    cin >> n >> m;
    for(LL i=1;i<=n;i++){
        cin >> a[i];
        if(vc[ a[i] ].empty()){
            vc[ a[i] ].push_back(0);
            b[++tot] = a[i];
        } 
        vc[ a[i] ].push_back(i);
    }
    LL dis = n+1;
    for(LL i=1;i<=tot;i++){
        LL t = b[i]; vc[t].push_back(n+1);
        LL l, r = -1, len = vc[t].size();
        for(LL j=1;j<len;j++){
            l = r+1;
            r = l-(vc[t][j]-vc[t][j-1]-1);
            LL L = r+dis, R = l+dis;
            ans += query(R-1) - (L-2>0? query(L-2):0);
            add(L, 1); add(R+1, -1);
        }
        r = -1;
        for(LL j=1;j<len;j++){
            l = r+1;
            r = l-(vc[t][j]-vc[t][j-1]-1);
            LL L = r+dis, R = l+dis;
            add(L, -1); add(R+1, 1);
        }
    }
    cout << ans << "\n";

    clear();

}

int main(){
    ios::sync_with_stdio(false);
    LL ce;
    ce = 1;
    //cin >> ce;
    while(ce--){
        solve();
    }


    return 0;
}