// C++ Versionintt1[MAXN],t2[MAXN],n;inlineintlowbit(intx){returnx&(-x);}voidadd(intk,intv){intv1=k*v;while(k<=n){t1[k]+=v,t2[k]+=v1;k+=lowbit(k);}}intgetsum(int*t,intk){intret=0;while(k){ret+=t[k];k-=lowbit(k);}returnret;}voidadd1(intl,intr,intv){add(l,v),add(r+1,-v);// 将区间加差分为两个前缀加}longlonggetsum1(intl,intr){return(r+1ll)*getsum(t1,r)-1ll*l*getsum(t1,l-1)-(getsum(t2,r)-getsum(t2,l-1));}
O(\log n) 查询第 k 小/大元素。在此处只讨论第 k 小,第 k 大问题可以通过简单计算转化为第 k 小问题。
参考 "可持久化线段树" 章节中,关于求区间第 k 小的思想。将所有数字看成一个可重集合,即定义数组 a 表示值为 i 的元素在整个序列重出现了 a_i 次。找第 k 大就是找到最小的 x 恰好满足 \sum_{i=1}^{x}a_i \geq k
因此可以想到算法:如果已经找到 x 满足 \sum_{i=1}^{x}a_i \le k,考虑能不能让 x 继续增加,使其仍然满足这个条件。找到最大的 x 后,x+1 就是所要的值。
在树状数组中,节点是根据 2 的幂划分的,每次可以扩大 2 的幂的长度。令 sum 表示当前的 x 所代表的前缀和,有如下算法找到最大的 x:
求出 depth=\left \lfloor \log_2n \right \rfloor
计算 t=\sum_{i=x+1}^{x+2^{depth}}a_i
如果 sum+t \le k,则此时扩展成功,将 2^{depth} 累加到 x 上;否则扩展失败,对 x 不进行操作
将 depth 减 1,回到步骤 2,直至 depth 为 0
1 2 3 4 5 6 7 8 910111213
// C++ Version// 权值树状数组查询第k小intkth(intk){intcnt=0,ret=0;for(inti=log2(n);~i;--i){// i 与上文 depth 含义相同ret+=1<<i;// 尝试扩展if(ret>=n||cnt+t[ret]>=k)// 如果扩展失败ret-=1<<i;elsecnt+=t[ret];// 扩展成功后 要更新之前求和的值}returnret+1;}
1 2 3 4 5 6 7 8 9101112
# Python Version# 权值树状数组查询第 k 小defkth(k):cnt=0;ret=0i=log2(n)# i 与上文 depth 含义相同while~i:ret=ret+(1<<i)# 尝试扩展ifret>=norcnt+t[ret]>=k:# 如果扩展失败ret=ret-(1<<i)else:cnt=cnt+t[ret]# 扩展成功后 要更新之前求和的值returnret+1
时间戳优化:
对付多组数据很常见的技巧。如果每次输入新数据时,都暴力清空树状数组,就可能会造成超时。因此使用 tag 标记,存储当前节点上次使用时间(即最近一次是被第几组数据使用)。每次操作时判断这个位置 tag 中的时间和当前时间是否相同,就可以判断这个位置应该是 0 还是数组内的值。
1 2 3 4 5 6 7 8 910111213141516171819202122
// C++ Version// 时间戳优化inttag[MAXN],t[MAXN],Tag;voidreset(){++Tag;}voidadd(intk,intv){while(k<=n){if(tag[k]!=Tag)t[k]=0;t[k]+=v,tag[k]=Tag;k+=lowbit(k);}}intgetsum(intk){intret=0;while(k){if(tag[k]==Tag)ret+=t[k];k-=lowbit(k);}returnret;}