跳转至

KDTree

带插入区间求和

Dimen 为维度,根据题目要求维护的数量进行修改。

aplpha 为平衡因子,可以卡常适当修改。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
namespace KDTree
{
    constexpr int Dimen=2;
    constexpr double alpha=0.777;
    struct point{array<int,Dimen> pos; int val;}a[MAX],info[MAX];
    int lson[MAX],rson[MAX],Max[MAX][Dimen],Min[MAX][Dimen],sum[MAX],siz[MAX];
    int stk[MAX],top;
    int root,tot;
    #define nxtwd ((wd+1)%Dimen)
    inline int New() {return top?stk[top--]:++tot;}
    inline void pushup(int i)
    {
        siz[i]=siz[lson[i]]+siz[rson[i]]+1,sum[i]=sum[lson[i]]+sum[rson[i]]+info[i].val;
        for(int j=0;j<Dimen;++j)
        {
            Max[i][j]=Min[i][j]=info[i].pos[j];
            if(lson[i]) cmax(Max[i][j],Max[lson[i]][j]),cmin(Min[i][j],Min[lson[i]][j]);
            if(rson[i]) cmax(Max[i][j],Max[rson[i]][j]),cmin(Min[i][j],Min[rson[i]][j]);
        }
    }
    int build(int l,int r,int wd)
    {
        if(l>r) return 0;
        int mid=(l+r)>>1,k=New();
        nth_element(a+l,a+mid,a+r+1,[&](point x,point y){return x.pos[wd]<y.pos[wd];});
        return info[k]=a[mid],lson[k]=build(l,mid-1,nxtwd),rson[k]=build(mid+1,r,nxtwd),pushup(k),k;
    }
    void Work(int i,int num)
    {
        if(lson[i]) Work(lson[i],num);
        a[siz[lson[i]]+num+1]=info[i],stk[++top]=i;
        if(rson[i]) Work(rson[i],num+siz[lson[i]]+1);
    }
    inline void Rebuild(int &i,int wd) {if(siz[i]*alpha<siz[lson[i]]||siz[i]*alpha<siz[rson[i]]) Work(i,0),i=build(1,siz[i],wd);}
    void Ins(int &i,int wd,point k)
    {
        if(!i) return i=New(),lson[i]=rson[i]=0,info[i]=k,pushup(i);
        if(k.pos[wd]<=info[i].pos[wd]) Ins(lson[i],nxtwd,k); else Ins(rson[i],nxtwd,k);
        pushup(i),Rebuild(i,wd);
    }
    int Query(int i,array<int,Dimen> L,array<int,Dimen> R)
    {
        if(!i) return 0;
        bool fl1=1,fl2=0,fl3=1;
        for(int j=0;j<Dimen;++j) fl1&=(L[j]<=Min[i][j]&&Max[i][j]<=R[j]),fl2|=(Max[i][j]<L[j]||R[j]<Min[i][j]);
        if(fl1) return sum[i];
        if(fl2) return 0;
        for(int j=0;j<Dimen;++j) fl3&=(L[j]<=info[i].pos[j]&&info[i].pos[j]<=R[j]);
        return Query(lson[i],L,R)+Query(rson[i],L,R)+fl3*info[i].val;
    }
    #undef nxtwd
}