[leetcode刷题笔记]线段树Segment Tree

在做两道区间检索题中接触到了线段树的概念，本文将其进行整理，文末几篇参考对线段树做了系统的介绍，结合两道leetcode题，对线段树的创建、查找、更新有了更深地掌握。文中不足，还望多多指正。

区域和检索 - 数组不可变

给定一个整数数组 nums，求出数组从索引 i 到 j (i ≤ j) 范围内元素的总和，包含 i, j 两点。

给定 nums = [-2, 0, 3, -5, 2, -1]，求和函数为 sumRange()

sumRange(0, 2) -> 1
sumRange(2, 5) -> -1
sumRange(0, 5) -> -3

线段树(Segment Tree)也是一棵树，只不过元素的值代表一个区间。常用区间的统计操作，比如一个区间的最大值(max)，最小值(min)，和(sum)等等，线段树是一个平衡二叉树，但不一定是完全二叉树。

在本题中，数组nums=[-2, 0, 3, -5, 2, -1]构建一棵线段树如下，节点为区间和(sum)

根节点就是 0~lenght-1 的和，根节点的左右子树平分根节点的区间，然后依次类推，直到只有一个元素不能划分为止，该元素也就是二叉树的叶子节点。

求线段树的区间统计，时间复杂度和二叉树的高度有关系，和元素的个数没关系，它的时间复杂度为 O(log n)。

如果我们用数组来存储线段树的话，我们大致需要开辟多大的数组空间呢？
h层的满二叉树总共有 2^h-1 个节点，第h-1层有2^(h-1)个节点，它们大概是两倍的关系。也就是说对于满二叉树 最后一层的节点数乘以2 大致就是整棵树的节点数。但是线段树并不一定是满二叉树，但是一定是平衡二叉树，所以需要多冗余一层。也就是 乘以4 就足以盛放所有的节点数，但是会浪费一定的内存空间。

class SegmentTree:
    def __init__(self, data:List[int]): 
        '''
        data:传入的数组
        merge:处理的业务逻辑，例如求和/最大值/最小值，lambda表达式
        '''
        self.data = data
        self.n = len(data)
        #  申请4倍data长度的空间来存线段树节点
        self.tree = [None] * (4 * self.n) # 索引i的左孩子索引为2i+1，右孩子为2i+2
        if self.n:
            self.build(0, 0, self.n-1)

    def querySum(self, ql:int, qr:int) -> int:
        '''
        返回区间[ql,..,qr]的值
        '''
        return self.query(0, 0, self.n-1, ql, qr)

    def build(self, tree_index:int, l:int, r:int):
        '''
        递归创建线段树
        tree_index : 线段树节点在数组中位置
        l, r : 该节点表示的区间的左,右边界
        '''
        if l == r:
            self.tree[tree_index] = self.data[l]
            return
        mid = (l+r) // 2 # 区间中点,对应左孩子区间结束,右孩子区间开头
        left, right = 2 * tree_index + 1, 2 * tree_index + 2 # tree_index的左右子树索引
        self.build(left, l, mid)
        self.build(right, mid+1, r)
        self.tree[tree_index] = self.tree[left] + self.tree[right]

查找查找分为四种情况

• 查询的区间在刚好就是当前节点表示区间
  查询结果即为当前节点值

• 查询区间全在左子树

• 查询区间全在右子树

• 查询区间部分在左子树，部分在右子树
  需到左右子树分别查询，并将查询结果求和（本题构造求和的线段树）

   def query(self, tree_index:int, l:int, r:int, ql:int, qr:int) -> int:
       '''
       递归查询区间[ql,..,qr]的值
       tree_index : 某个根节点的索引
       l, r : 该节点表示的区间的左右边界
       ql, qr: 待查询区间的左右边界
       '''
       if l == ql and r == qr:
           return self.tree[tree_index]
  
       # 区间中点,对应左孩子区间结束,右孩子区间开头
       mid = (l+r) // 2 
       left, right = tree_index * 2 + 1, tree_index * 2 + 2
       if qr <= mid:
           # 查询区间全在左子树
           return self.query(left, l, mid, ql, qr)
       elif ql > mid:
           # 查询区间全在右子树
           return self.query(right, mid+1, r, ql, qr)
  
       # 查询区间一部分在左子树一部分在右子树
       return self.query(left, l, mid, ql, mid) + self.query(right, mid+1, r, mid+1, qr)
  

构建好线段树后，本题操作

class NumArray:

    def __init__(self, nums: List[int]):
        self.segment_tree = SegmentTree(nums)
    

    def sumRange(self, i: int, j: int) -> int:
        return self.segment_tree. querySum(i, j)

区域和检索 - 数组可修改

给定一个整数数组 nums，求出数组从索引 i 到 j (i ≤ j) 范围内元素的总和，包含 i, j 两点。

update(i, val) 函数可以通过将下标为 i 的数值更新为 val，从而对数列进行修改。

Given nums = [1, 3, 5]

sumRange(0, 2) -> 9
update(1, 2)
sumRange(0, 2) -> 8

数组仅可以在 update 函数下进行修改。
你可以假设 update 函数与 sumRange 函数的调用次数是均匀分布的。

相比上一题增加了update操作，其更新操作需找到相应的叶子结点，将其值更新，然后将包括其大区间的值更新。

def update(self, tree_index, l, r, index):
    '''
    tree_index:某个根节点索引
    l, r : 此根节点代表区间的左右边界
    index : 更新的值的索引
    '''
    if l == r==index :
        self.tree[tree_index] = self.data[index]
        return
    mid = (l+r)//2
    left, right = 2 * tree_index + 1, 2 * tree_index + 2
    if index > mid:
        # 要更新的区间在右子树
        self.update(right, mid+1, r, index)
    else:
        # 要更新的区间在左子树index<=mid
        self.update(left, l, mid, index)

    self.tree[tree_index] = self.tree[left] + self.tree[right]

def updateSum(self,index:int,value:int):
    self.data[index] = value
    self.update(0, 0, self.n-1, index)

Reference

1.数据结构与算法（十）线段树(Segment Tree)入门
2.线段树SegmentTree-数组实现