HDU 3854 Glorious Array(树状数组)

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
#include<time.h>
#include<iostream>
#define INF 10000000
#define LL long long
using namespace std;
int n,K;
vector<int> vec;
int col[100005];
struct BIT
{
    int dat[100005];
    void clear()
    {
        memset(dat,0,sizeof(dat));
    }
    int lowBit(int x)
    {
        return -x&x;
    }
    int getSum(int x)
    {
        int s=0;
        while(x>0)
        {
            s+=dat[x];
            x-=lowBit(x);
        }
        return s;
    }
    void addVal(int x,int c)
    {
        if(x<=0) return;
        while(x<=n)
        {
            dat[x]+=c;
            x+=lowBit(x);
        }
    }
    void setVal(int x,int c)
    {
        int old=getSum(x)-getSum(x-1);
        addVal(x,-old+c);
    }
    int getCol(int x)
    {
        return getSum(x)-getSum(x-1);
    }
    int countBlack(int ql,int qr)
    {
        if(qr<ql) return 0;
        return getSum(qr)-getSum(ql-1);
    }
    int countWhite(int ql,int qr)
    {
        if(qr<ql) return 0;
        int s=countBlack(ql,qr);
        return qr-ql+1-s;
    }
};
BIT tree;
int leftM[100005],rightM[100005];
int main()
{
    int T;
    scanf("%d",&T);
    while(T--)
    {
        int m;
        scanf("%d%d%d",&n,&m,&K);
        vec.clear();
        for(int i=1; i<=n; ++i)
        {
            int val;
            scanf("%d",&val);
            if(val<K)
                vec.push_back(i);
        }
        tree.clear();
        for(int i=1; i<=n; ++i)
        {
            int val;
            scanf("%d",&val);
            tree.setVal(i,val);
            int l=upper_bound(vec.begin(),vec.end(),i)-vec.begin()-1;
            int r=lower_bound(vec.begin(),vec.end(),i)-vec.begin();
            if(l<0) leftM[i]=-1;
            else leftM[i]=vec[l];
            if(r>=vec.size()) rightM[i]=-1;
            else rightM[i]=vec[r];
        }
        LL sum=0;
        int last=1,end;
        for(int i=0; i<vec.size(); ++i)
        {
            last=(i==0)?1:vec[i-1]+1;
            end=(i==vec.size()-1)?n:vec[i+1];
            LL lb=tree.countBlack(last,vec[i]-1),rw=tree.countWhite(vec[i]+1,n);
            LL lw=tree.countWhite(last,vec[i]-1),rb=tree.countBlack(vec[i]+1,n);
            sum+=lb*rw;
            sum+=lw*rb;
            int col=tree.getCol(vec[i]);
            if(col==0)
                sum+=lb+rb;
            else
                sum+=lw+rw;
        }
        while(m--)
        {
            int p;
            scanf("%d",&p);
            if(p==1) printf("%I64d\n",sum);
            else
            {
                int x;
                scanf("%d",&x);
                int old=tree.getCol(x);
                if(old==1)
                {
                    if(leftM[x]!=-1)
                    {
                        LL lw,lb;
                        if(leftM[x]==x)
                        {
                            lw=tree.countWhite(1,leftM[x]-1);
                            lb=tree.countBlack(1,leftM[x]-1);
                        }
                        else
                        {
                            lw=tree.countWhite(1,leftM[x]);
                            lb=tree.countBlack(1,leftM[x]);
                        }
                        sum-=lw;
                        sum+=lb;
                    }
                    if(rightM[x]!=-1)
                    {
                        LL rw,rb;
                        if(rightM[x]==x)
                        {
                            rb=tree.countBlack(rightM[x]+1,n);
                            rw=tree.countWhite(rightM[x]+1,n);
                        }
                        else
                        {
                            rb=tree.countBlack(rightM[x],n);
                            rw=tree.countWhite(rightM[x],n);
                        }
                        sum-=rw;
                        sum+=rb;
                    }
                }
                else
                {
                    if(leftM[x]!=-1)
                    {
                        LL lw,lb;
                        if(leftM[x]==x)
                        {
                            lw=tree.countWhite(1,leftM[x]-1);
                            lb=tree.countBlack(1,leftM[x]-1);
                        }
                        else
                        {
                            lw=tree.countWhite(1,leftM[x]);
                            lb=tree.countBlack(1,leftM[x]);
                        }
                        sum+=lw;
                        sum-=lb;
                    }
                    if(rightM[x]!=-1)
                    {
                        LL rw,rb;
                        if(rightM[x]==x)
                        {
                            rb=tree.countBlack(rightM[x]+1,n);
                            rw=tree.countWhite(rightM[x]+1,n);
                        }
                        else
                        {
                            rb=tree.countBlack(rightM[x],n);
                            rw=tree.countWhite(rightM[x],n);
                        }
                        sum+=rw;
                        sum-=rb;
                    }
                }
                tree.setVal(x,old^1);
            }
        }
    }
    return 0;
}