洛谷3372-线段树1-python-(线段树)

看注释

global n,m,nums,MAX_LENS,tree

MAX_LENS=10000
#树存储结构
tree=[0]*MAX_LENS
#懒惰标记存储结构
lazy_tag=[0]*MAX_LENS


#建树
def build_tree(node,sta,end):
    if sta==end:
        tree[node]=nums[sta]
        return
    mid=(sta+end)//2
    left_node=2*node+1
    right_node=2*node+2
    build_tree(left_node,sta,mid)
    build_tree(right_node,mid+1,end)
    tree[node]=tree[left_node]+tree[right_node]
    
#查询
def query_tree(node,sta,end,L,R):
    if R<sta or L>end:
        return 0
    elif L<=sta and R>=end:
        return tree[node]
    else:
        mid=(sta+end)//2
        left_node=2*node+1
        right_node=2*node+2
        
        return query_tree(left_node,sta,mid,L,R) + query_tree(right_node,mid+1,end,L,R)

#单点更新
def modify_tree_onepoint(node,sta,end,index,value):
    if sta==end:
        nums[index]=value
        tree[node]=value
        return
    
    mid=(sta+end)//2
    left_node=node*2+1
    right_node=node*2+2
    if sta<=index<=mid:
        modify_tree_onepoint(left_node,sta,mid,index,value)
    else:
        modify_tree_onepoint(right_node,mid+1,end,index,value)
    tree[node]=tree[left_node]+tree[right_node]
 
#区间更新 引入lazy_tag数组
def modify_tree_interval(node,sta,end,c,l,r):
    
    if l<=sta and r>=end:
        tree[node]+=(end-sta+1)*c
        lazy_tag[node]+=c
        return
    mid=(sta+end)//2
    #已被标记
    if lazy_tag[node]:
        lazy_tag[2*node+1]+=lazy_tag[node]
        lazy_tag[2*node+2]+=lazy_tag[node]
        
        tree[2*node+1]+=(mid-sta+1)*lazy_tag[node]
        tree[2*node+2]+=(end-mid)*lazy_tag[node]

        lazy_tag[node]=0
    #左边或者右边有交集
    if l<=mid:
        modify_tree_interval(2*node+1,sta,mid,c,l,r)
    if r>mid:
        modify_tree_interval(2*node+2,mid+1,end,c,l,r)
    tree[node]=tree[2*node+1]+tree[2*node+2]

#带有懒惰标记的区间查询
def query_interval(node,sta,end,L,R):
    if L<=sta and R>=end:
        return tree[node]
    left_node=2*node+1
    right_node=2*node+2
    mid=(sta+end)//2
    if lazy_tag[node]:
        lazy_tag[left_node]+=lazy_tag[node]
        lazy_tag[right_node]+=lazy_tag[node]

        tree[left_node]+=(mid-sta+1)*lazy_tag[node]
        tree[right_node]+=(end-mid)*lazy_tag[node]

        lazy_tag[node]=0
    if L<=mid and R>mid:
        return query_interval(left_node,sta,mid,L,R)+query_interval(right_node,mid+1,end,L,R)
    elif L<=mid:
        return query_interval(left_node,sta,mid,L,R)
    else:
        return query_interval(right_node,mid+1,end,L,R)
    
#打印线段树
def print_tree():
    for i in range(12):
        print("node[%d]=%d"%(i,tree[i]))
    print()
    
#打印懒惰标记
def print_lazy_tag():
    for i in range(12):
        print("lazy_tag[%d]=%d"%(i,lazy_tag[i]))
    print()
    
n,m=map(int,input().split())
nums=list(map(int,input().split()))
"""
5 5
1 5 4 2 3
"""

build_tree(0,0,n-1)
print_tree()

print("2,3区间的和为：",query_tree(0,0,n-1,2,3))
print()

print("将下标为3的元素改为99后：")
modify_tree_onepoint(0,0,n-1,3,99)
print(nums)
print_tree()

modify_tree_interval(0,0,n-1,1,3,4)
print_tree()
print_lazy_tag()

#查询3-4
print(query_interval(0,0,n-1,2,4))
print_tree()
print_lazy_tag()