test/library-checker/data-structure/point_add_range_sum.test.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00
• Problem: https://judge.yosupo.jp/problem/point_add_range_sum

Depends on

• cp_library/ds/tree/bit/bit_cls.py

Code

# verification-helper: PROBLEM https://judge.yosupo.jp/problem/point_add_range_sum
import io,os
from cp_library.ds.tree.bit.bit_cls import BIT

input = io.BytesIO(os.read(0,os.fstat(0).st_size)).readline
MI = lambda : map(int, input().split())

n,q = MI()
a = [int(s) for s in input().split()]
# b = a.copy()
bit = BIT(a)
# for i in range(n):
#     assert b[i] == bit[i], f"{a[i]} != {bit[i]}"
ans = []
for i in range(q):
    t,p,x = MI()
    if t: ans.append(bit.range_sum(p,x))
    else: bit.add(p,x)

os.write(1," ".join(map(str,ans)).encode())

# verification-helper: PROBLEM https://judge.yosupo.jp/problem/point_add_range_sum
import io,os
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from typing import Union


'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
            ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓            
            ┃                                    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: Union[int, list[int]]):
        if isinstance(v, int): bit.d, bit.n = [0]*v, v
        else: bit.build(v)
        bit.lb = 1<<(bit.n.bit_length()-1)

    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 range_sum(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 0
    
    def bisect_right(bit, v) -> int:
        i = s = 0; ni = m = bit.lb
        while m:
            if ni <= bit.n and (ns:=s+bit.d[ni-1]) <= v: s, i = ns, ni
            ni = (m:=m>>1)|i
        return i

input = io.BytesIO(os.read(0,os.fstat(0).st_size)).readline
MI = lambda : map(int, input().split())

n,q = MI()
a = [int(s) for s in input().split()]
# b = a.copy()
bit = BIT(a)
# for i in range(n):
#     assert b[i] == bit[i], f"{a[i]} != {bit[i]}"
ans = []
for i in range(q):
    t,p,x = MI()
    if t: ans.append(bit.range_sum(p,x))
    else: bit.add(p,x)

os.write(1," ".join(map(str,ans)).encode())

Back to top page