cp-library
This documentation is automatically generated by online-judge-tools/verification-helper
test/library-checker/data-structure/rectangle_sum_wm_bit.test.py
- View this file on GitHub
- Last update: 2025-06-20 03:24:59+09:00
- Problem: https://judge.yosupo.jp/problem/rectangle_sum
Depends on
- cp_library/alg/dp/max2_fn.py
- cp_library/bit/pack/pack_sm_fn.py
- cp_library/bit/popcnt32_fn.py
- cp_library/ds/array/u32f_fn.py
- cp_library/ds/tree/bit/bit_cls.py
- cp_library/ds/wavelet/bit_array_cls.py
- cp_library/ds/wavelet/wm_bit_cls.py
- cp_library/ds/wavelet/wm_static_cls.py
- cp_library/ds/wavelet/wm_weighted_cls.py
Code
# verification-helper: PROBLEM https://judge.yosupo.jp/problem/rectangle_sum
def main():
N, Q = map(int, input().split())
Xn, Yn, Wn = [0]*N, [0]*N, [0]*N; L, D, R, U = [0]*Q, [0]*Q, [0]*Q, [0]*Q
for i in range(N): Xn[i], Yn[i], Wn[i] = map(int, input().split())
for i in range(Q): L[i], D[i], R[i], U[i] = map(int, input().split())
icoord_compress_multi(L, R, Xn, distinct=True); icoord_compress_multi(D, U, Yn)
Y, W = [0]*N, [0]*N
for i,j in enumerate(Xn): Y[j], W[j] = Yn[i], Wn[i]
wm = WMBIT(Y, W)
for i in range(Q): append(str(wm.sum_rect(L[i], D[i], R[i], U[i]))); append('\n')
os.write(1, sb.build().encode())
from cp_library.ds.wavelet.wm_bit_cls import WMBIT
from cp_library.alg.dp.max2_fn import max2
from cp_library.bit.pack.pack_sm_fn import pack_sm
def icoord_compress_multi(*A: list[int], distinct=False):
N = mx = 0
for Ai in A: N += len(Ai); mx = max2(mx, len(Ai))
si, mi = pack_sm(mx-1); sj, mj = pack_sm((len(A)-1)<<si)
S, k = [0]*N, 0
for i,Ai in enumerate(A):
for j,a in enumerate(Ai): S[k]=a << sj | i << si | j; k += 1
S.sort(); r = p = -1
for aji in S:
a, i, j = aji >> sj, (aji&mj) >> si , aji & mi
if i == 2 and (distinct or a != p): r = r+1; p = a
A[i][j] = r+(i!=2)
return A
import sys,os
from __pypy__ import builders # type: ignore
sb = builders.StringBuilder()
append = sb.append
def input(): return sys.stdin.buffer.readline().strip()
if __name__ == "__main__":
main()# verification-helper: PROBLEM https://judge.yosupo.jp/problem/rectangle_sum
def main():
N, Q = map(int, input().split())
Xn, Yn, Wn = [0]*N, [0]*N, [0]*N; L, D, R, U = [0]*Q, [0]*Q, [0]*Q, [0]*Q
for i in range(N): Xn[i], Yn[i], Wn[i] = map(int, input().split())
for i in range(Q): L[i], D[i], R[i], U[i] = map(int, input().split())
icoord_compress_multi(L, R, Xn, distinct=True); icoord_compress_multi(D, U, Yn)
Y, W = [0]*N, [0]*N
for i,j in enumerate(Xn): Y[j], W[j] = Yn[i], Wn[i]
wm = WMBIT(Y, W)
for i in range(Q): append(str(wm.sum_rect(L[i], D[i], R[i], U[i]))); append('\n')
os.write(1, sb.build().encode())
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
https://kobejean.github.io/cp-library
'''
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ 7 ┃
┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━┛
┏━━━━━━━━━━━━━━━━━━┓ │
┃ 3 ┃◄────────────────┤
┗━━━━━━━━━━━━━━━━┯━┛ │
┏━━━━━━━━┓ │ ┏━━━━━━━━┓ │
┃ 1 ┃◄──────┤ ┃ 5 ┃◄──────┤
┗━━━━━━┯━┛ │ ┗━━━━━━┯━┛ │
┏━━━┓ │ ┏━━━┓ │ ┏━━━┓ │ ┏━━━┓ │
┃ 0 ┃◄─┤ ┃ 2 ┃◄─┤ ┃ 4 ┃◄─┤ ┃ 6 ┃◄─┤
┗━┯━┛ │ ┗━┯━┛ │ ┗━┯━┛ │ ┗━┯━┛ │
│ │ │ │ │ │ │ │
▼ ▼ ▼ ▼ ▼ ▼ ▼ ▼
┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓
┃ 0 ┃┃ 1 ┃┃ 2 ┃┃ 3 ┃┃ 4 ┃┃ 5 ┃┃ 6 ┃┃ 7 ┃
┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
Data Structure - Tree - Binary Index Tree
'''
class BIT:
def __init__(bit, v):
if isinstance(v, int): bit._d, bit._n = [0]*v, v
else: bit.build(v)
bit._lb = 1<<bit._n.bit_length()
def build(bit, data):
bit._d, bit._n = data, len(data)
for i in range(bit._n):
if (r := i|i+1) < bit._n: bit._d[r] += bit._d[i]
def add(bit, i, x):
while i < bit._n: bit._d[i] += x; i |= i+1
def sum(bit, n: int) -> int:
s = 0
while n: s, n = s+bit._d[n-1], n&n-1
return s
def sum_range(bit, l, r):
s = 0
while r: s, r = s+bit._d[r-1], r&r-1
while l: s, l = s-bit._d[l-1], l&l-1
return s
def __len__(bit) -> int:
return bit._n
def __getitem__(bit, i: int) -> int:
s, l = bit._d[i], i&(i+1)
while l != i: s, i = s-bit._d[i-1], i-(i&-i)
return s
get = __getitem__
def __setitem__(bit, i: int, x: int) -> None:
bit.add(i, x-bit[i])
set = __setitem__
def prelist(bit) -> list[int]:
pre = [0]+bit._d
for i in range(bit._n+1): pre[i] += pre[i&i-1]
return pre
def bisect_left(bit, v) -> int:
return bit.bisect_right(v-1) if v>0 else -1
def bisect_right(bit, v, key=None) -> int:
i = s = 0; m = bit._lb
if key:
while m := m>>1:
if (ni := m|i) <= bit._n and key(ns:=s+bit._d[ni-1]) <= v: s, i = ns, ni
else:
while m := m>>1:
if (ni := m|i) <= bit._n and (ns:=s+bit._d[ni-1]) <= v: s, i = ns, ni
return i
class BitArray:
def __init__(B, N: int):
B.N, B.Z = N, (N+31)>>5
B.bits, B.cnt = u32f(B.Z+1), u32f(B.Z+1)
def build(B):
B.bits.pop()
for i,b in enumerate(B.bits): B.cnt[i+1] = B.cnt[i]+popcnt32(b)
B.bits.append(1)
def __len__(B): return B.N
def __getitem__(B, i: int): return B.bits[i>>5]>>(31-(i&31))&1
def set0(B, i: int): B.bits[i>>5]&=~(1<<31-(i&31))
def set1(B, i: int): B.bits[i>>5]|=1<<31-(i&31)
def count0(B, r: int): return r-B.count1(r)
def count1(B, r: int): return B.cnt[r>>5]+popcnt32(B.bits[r>>5]>>32-(r&31))
def select0(B, k: int):
if not 0<=k<B.N-B.cnt[-1]: return -1
l,r,k=0,B.N,k+1
while 1<r-l:
if B.count0(m:=(l+r)>>1)<k:l=m
else:r=m
return l
def select1(B, k: int):
if not 0<=k<B.cnt[-1]: return -1
l,r,k=0,B.N,k+1
while 1<r-l:
if B.count1(m:=(l+r)>>1)<k:l=m
else:r=m
return l
def popcnt32(x):
x = ((x >> 1) & 0x55555555) + (x & 0x55555555)
x = ((x >> 2) & 0x33333333) + (x & 0x33333333)
x = ((x >> 4) & 0x0f0f0f0f) + (x & 0x0f0f0f0f)
x = ((x >> 8) & 0x00ff00ff) + (x & 0x00ff00ff)
x = ((x >> 16) & 0x0000ffff) + (x & 0x0000ffff)
return x
if hasattr(int, 'bit_count'):
popcnt32 = int.bit_count
from array import array
def u32f(N: int, elm: int = 0): return array('I', (elm,))*N # unsigned int
class WMStatic:
class Level(BitArray):
def __init__(L, N: int, H: int):
super().__init__(N)
L.H = H
def build(L):
super().build()
L.T0, L.T1 = L.N-L.cnt[-1], L.cnt[-1]
def pos(L, bit: int, i: int): return L.T0+L.count1(i) if bit else L.count0(i)
def pos2(L, bit: int, i: int, j: int): return (L.T0+L.count1(i), L.T0+L.count1(j)) if bit else (L.count0(i), L.count0(j))
def __init__(wm,A,Amax:int=None):wm._build(A,[0]*len(A),max(A,default=0)if Amax is None else Amax)
def _build(wm, A, nA, Amax):wm.N,wm.H=len(A),Amax.bit_length();wm._build_levels(A,nA)
def _build_levels(wm, A, nA):
wm.up=[wm.Level(wm.N,H) for H in range(wm.H)];wm.down=wm.up[::-1]
for L in wm.down:
x,y,i=-1,wm.N-1,wm.N
while i:y-=A[i:=i-1]>>L.H&1
for i,a in enumerate(A):
if a>>L.H&1:nA[y:=y+1]=a;L.set1(i)
else:nA[x:=x+1]=a
A,nA=nA,A;L.build()
def __getitem__(wm,i):
y=0
for L in wm.down:y=y<<1|(bit:=L[i]);i=L.pos(bit,i)
return y
def kth(wm, k: int, l: int, r: int):
'''Returns the `k+1`-th value in sorted order of values in range `[l, r)`'''
s=0
for L in wm.down:
l,r=l-(l1:=L.count1(l)),r-(r1:=L.count1(r))
if k>=r-l:s|=1<<L.H;k-=r-l;l,r=L.T0+l1,L.T0+r1
return s
def select(wm, y: int, k: int, l: int = 0, r: int = -1):
'''Returns the index of the `k+1`-th occurance of `y` in range `[l, r)`'''
if not(0<=y<1<<wm.H):return-1
if r==-1:r=wm.N-1
for L in wm.down:l,r=L.pos2(L[y],l,r)
if not l<=(i:=l+k)<r:return-1
for L in wm.up:
if y>>L.H&1:i=L.select1(i-L.T0)
else:i=L.select0(i)
return i
def rank(wm, y: int, r: int): return wm.rank_range(y, 0, r)
def rank_range(wm, y: int, l: int, r: int):
if l >= r: return 0
for L in wm.down:l,r=L.pos2(L[y],l,r)
return r-l
def count_at(wm, y: int, l: int, r: int):
'''Count how many `y` values are in range `[l,r)` '''
if l >= r: return 0
return wm._cnt(y+1, l, r)-wm._cnt(y, l, r)
def count_below(wm, u: int, l: int, r: int):
'''Count `i`'s in `[l,r)` such that `A[i] < u` '''
return wm._cnt(u, l, r)
def count_between(wm, d: int, u: int, l: int, r: int):
'''Count `i`'s in `[l,r)` such that `d <= A[i] < u` '''
if l >= r or d >= u: return 0
return wm._cnt(u, l, r)-wm._cnt(d, l, r)
def _cnt(wm, u: int, l: int, r: int):
if u<=0:return 0
if wm.H<u.bit_length():return r-l
cnt=0
for L in wm.down:
l,r=l-(l1:=L.count1(l)),r-(r1:=L.count1(r))
if u>>L.H&1:cnt+=r-l;l,r=L.T0+l1,L.T0+r1
return cnt
def prev_val(wm,u:int,l:int,r:int):return wm.kth(cnt-1, l, r)if(cnt:=wm._cnt(u,l,r))else-1
def next_val(wm,d:int,l:int,r:int):return wm.kth(cnt, l, r)if(cnt:=wm._cnt(d,l,r))<r-l else-1
class WMWeighted(WMStatic):
class Level(WMStatic.Level):
def __init__(L,N:int,H:int):super().__init__(N,H);L.W=[0]*(N+1)
def build(L,W:list[int]):
super().build()
for i,w in enumerate(W):L.W[i+1]=L.W[i]+w
def sum(L,l:int,r:int):return L.W[r]-L.W[l]
def __init__(wm, A:list[int],W:list[int],Amax:int=None):wm._build(A,W,[0]*len(A),[0]*len(A),max(A,default=0)if Amax is None else Amax)
def _build(wm,A,W,nA,nW,Amax):wm.N,wm.H=len(A),Amax.bit_length();wm._build_base(W);wm._build_levels(A,W,nA,nW)
def _build_base(wm, W):
wm.W = [0]*(wm.N+1)
for i,w in enumerate(W): wm.W[i+1] = wm.W[i]+w
def _build_levels(wm, A, W, nA, nW):
wm.up = [wm.Level(wm.N, H) for H in range(wm.H)]; wm.down = wm.up[::-1]
for L in wm.down:
x,y,i=-1,wm.N-1,wm.N
while i:y-=A[i:=i-1]>>L.H&1
for i,a in enumerate(A):
if a>>L.H&1:nA[y:=y+1],nW[y]=a,W[i];L.set1(i)
else:nA[x:=x+1],nW[x]=a,W[i]
A,nA,W,nW=nA,A,nW,W;L.build(W)
def sum_range(wm,l:int,r:int):return wm._sum_range(l,r) if l<r else 0
def sum_at(wm,y:int,l:int,r:int):return wm._sum_rect(y,y+1,l,r) if l<r else 0
def sum_below(wm,u:int,l:int,r:int):return wm._sum_below(u,l,r) if l<r else 0
def sum_above(wm,d:int,l:int,r:int):return wm._sum_above(d,l,r) if l<r else 0
def sum_between(wm,d:int,u:int,l:int,r:int):return wm._sum_rect(d,u,l,r)if l<r and d<u else 0
def sum_corner(wm,r:int,u:int):return wm.sum_below(u,0,r) if 0<r else 0
def sum_rect(wm,l:int,d:int,r:int,u:int):return wm._sum_rect(d,u,l,r)if l<r and d<u else 0
def _sum_range(wm,l,r):return wm.W[r]-wm.W[l]
def _sum_below(wm,u,l,r):
if u<=0:return 0
elif wm.H<u.bit_length():return wm._sum_range(l,r)
sum = 0
for L in wm.down:
l,r=l-(l1:=L.count1(l)),r-(r1:=L.count1(r))
if u>>L.H&1:sum+=L.sum(l,r);l,r=L.T0+l1,L.T0+r1
return sum
def _sum_above(wm,d,l,r):
if d<=0:return wm._sum_range(l,r)
elif wm.H<d.bit_length():return 0
sum,d=0,d-1
for L in wm.down:
l0,r0=l-(l:=L.T0+L.count1(l)),r-(r:=L.T0+L.count1(r))
if d>>L.H&1==0:sum+=L.sum(l,r);l,r=L.T0+l0,L.T0+r0
return sum
def _sum_rect(wm,d,u,l,r):
if u<=0 or wm.H<d.bit_length():return 0
elif d<=0:return wm._sum_below(u,l,r)
elif wm.H<u.bit_length():return wm._sum_above(d,l,r)
same,sum,d=1,0,d-1
for L in wm.down:
db,ub,l,r=d>>L.H&1,u>>L.H&1,l-(l1:=L.count1(l)),r-(r1:=L.count1(r))
if same:
if db!=ub:same,dl,dr,l,r=0,l,r,L.T0+l1,L.T0+r1
elif db:l,r=L.T0+l1,L.T0+r1
else:
if ub:sum+=L.sum(l,r);l,r=L.T0+l1,L.T0+r1
dl0,dr0=dl-(dl:=L.T0+L.count1(dl)),dr-(dr:=L.T0+L.count1(dr))
if not db:sum+=L.sum(dl,dr);dl,dr=L.T0+dl0,L.T0+dr0
return sum
class WMBIT(WMWeighted):
class Level(WMStatic.Level):
def build(L,W:list[int]):super().build();L.W=BIT(W[:])
def sum(L,l:int,r:int):return L.W.sum_range(l,r)
def _build_base(wm,W):wm.W=BIT(W[:])
def _sum_range(wm,l,r):return wm.W.sum_range(l,r)
def add(wm,i:int,w:int):
wm.W.add(i,w)
for L in wm.down:L.W.add(i:=L.pos(L[i],i),w)
def max2(a, b):
return a if a > b else b
def pack_sm(N: int): s=N.bit_length(); return s,(1<<s)-1
def icoord_compress_multi(*A: list[int], distinct=False):
N = mx = 0
for Ai in A: N += len(Ai); mx = max2(mx, len(Ai))
si, mi = pack_sm(mx-1); sj, mj = pack_sm((len(A)-1)<<si)
S, k = [0]*N, 0
for i,Ai in enumerate(A):
for j,a in enumerate(Ai): S[k]=a << sj | i << si | j; k += 1
S.sort(); r = p = -1
for aji in S:
a, i, j = aji >> sj, (aji&mj) >> si , aji & mi
if i == 2 and (distinct or a != p): r = r+1; p = a
A[i][j] = r+(i!=2)
return A
import sys,os
from __pypy__ import builders # type: ignore
sb = builders.StringBuilder()
append = sb.append
def input(): return sys.stdin.buffer.readline().strip()
if __name__ == "__main__":
main()