在做两道区间检索题中接触到了线段树的概念,本文将其进行整理,文末几篇参考对线段树做了系统的介绍,结合两道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)