0370. 区间加法
大约 5 分钟
0370. 区间加法
- 标签：数组、前缀和
- 难度：中等
题目链接
题目大意
描述：给定一个数组的长度 length ，初始情况下数组中所有数字均为 0。再给定 k 个更新操作。其中每个操作是一个三元组 [startIndex, endIndex, inc]，表示将子数组 nums[startIndex ... endIndex] （包括 startIndex、endIndex）上所有元素增加 inc。
要求：返回 k 次操作后的数组。
示例：
- 示例 1：
给定 length = 5，即 nums = [0, 0, 0, 0, 0]
操作 [1, 3, 2] -> [0, 2, 2, 2, 0]
操作 [2, 4, 3] -> [0, 2, 5, 5, 3]
操作 [0, 2, -2] -> [-2, 0, 3, 5, 3]
解题思路
思路 1：线段树
- 初始化一个长度为
length，值全为0的nums数组。 - 然后根据
nums数组构建一棵线段树。每个线段树的节点类存储当前区间的左右边界和该区间的和。并且线段树使用延迟标记。 - 然后遍历三元组操作，进行区间累加运算。
- 最后从线段树中查询数组所有元素，返回该数组即可。
这样构建线段树的时间复杂度为 ，单次区间更新的时间复杂度为 ，单次区间查询的时间复杂度为 。总体时间复杂度为 。
思路 1：线段树代码
# 线段树的节点类
class SegTreeNode:
def __init__(self, val=0):
self.left = -1 # 区间左边界
self.right = -1 # 区间右边界
self.val = val # 节点值（区间值）
self.lazy_tag = None # 区间和问题的延迟更新标记
# 线段树类
class SegmentTree:
# 初始化线段树接口
def __init__(self, nums, function):
self.size = len(nums)
self.tree = [SegTreeNode() for _ in range(4 * self.size)] # 维护 SegTreeNode 数组
self.nums = nums # 原始数据
self.function = function # function 是一个函数，左右区间的聚合方法
if self.size > 0:
self.__build(0, 0, self.size - 1)
# 单点更新接口：将 nums[i] 更改为 val
def update_point(self, i, val):
self.nums[i] = val
self.__update_point(i, val, 0)
# 区间更新接口：将区间为 [q_left, q_right] 上的所有元素值加上 val
def update_interval(self, q_left, q_right, val):
self.__update_interval(q_left, q_right, val, 0)
# 区间查询接口：查询区间为 [q_left, q_right] 的区间值
def query_interval(self, q_left, q_right):
return self.__query_interval(q_left, q_right, 0)
# 获取 nums 数组接口：返回 nums 数组
def get_nums(self):
for i in range(self.size):
self.nums[i] = self.query_interval(i, i)
return self.nums
# 以下为内部实现方法
# 构建线段树实现方法：节点的存储下标为 index，节点的区间为 [left, right]
def __build(self, index, left, right):
self.tree[index].left = left
self.tree[index].right = right
if left == right: # 叶子节点，节点值为对应位置的元素值
self.tree[index].val = self.nums[left]
return
mid = left + (right - left) // 2 # 左右节点划分点
left_index = index * 2 + 1 # 左子节点的存储下标
right_index = index * 2 + 2 # 右子节点的存储下标
self.__build(left_index, left, mid) # 递归创建左子树
self.__build(right_index, mid + 1, right) # 递归创建右子树
self.__pushup(index) # 向上更新节点的区间值
# 单点更新实现方法：将 nums[i] 更改为 val，节点的存储下标为 index
def __update_point(self, i, val, index):
left = self.tree[index].left
right = self.tree[index].right
if left == right:
self.tree[index].val = val # 叶子节点，节点值修改为 val
return
mid = left + (right - left) // 2 # 左右节点划分点
left_index = index * 2 + 1 # 左子节点的存储下标
right_index = index * 2 + 2 # 右子节点的存储下标
if i <= mid: # 在左子树中更新节点值
self.__update_point(i, val, left_index)
else: # 在右子树中更新节点值
self.__update_point(i, val, right_index)
self.__pushup(index) # 向上更新节点的区间值
# 区间更新实现方法
def __update_interval(self, q_left, q_right, val, index):
left = self.tree[index].left
right = self.tree[index].right
if left >= q_left and right <= q_right: # 节点所在区间被 [q_left, q_right] 所覆盖
if self.tree[index].lazy_tag is not None:
self.tree[index].lazy_tag += val # 将当前节点的延迟标记增加 val
else:
self.tree[index].lazy_tag = val # 将当前节点的延迟标记增加 val
interval_size = (right - left + 1) # 当前节点所在区间大小
self.tree[index].val += val * interval_size # 当前节点所在区间每个元素值增加 val
return
if right < q_left or left > q_right: # 节点所在区间与 [q_left, q_right] 无关
return
self.__pushdown(index) # 向下更新节点的区间值
mid = left + (right - left) // 2 # 左右节点划分点
left_index = index * 2 + 1 # 左子节点的存储下标
right_index = index * 2 + 2 # 右子节点的存储下标
if q_left <= mid: # 在左子树中更新区间值
self.__update_interval(q_left, q_right, val, left_index)
if q_right > mid: # 在右子树中更新区间值
self.__update_interval(q_left, q_right, val, right_index)
self.__pushup(index) # 向上更新节点的区间值
# 区间查询实现方法：在线段树中搜索区间为 [q_left, q_right] 的区间值
def __query_interval(self, q_left, q_right, index):
left = self.tree[index].left
right = self.tree[index].right
if left >= q_left and right <= q_right: # 节点所在区间被 [q_left, q_right] 所覆盖
return self.tree[index].val # 直接返回节点值
if right < q_left or left > q_right: # 节点所在区间与 [q_left, q_right] 无关
return 0
self.__pushdown(index)
mid = left + (right - left) // 2 # 左右节点划分点
left_index = index * 2 + 1 # 左子节点的存储下标
right_index = index * 2 + 2 # 右子节点的存储下标
res_left = 0 # 左子树查询结果
res_right = 0 # 右子树查询结果
if q_left <= mid: # 在左子树中查询
res_left = self.__query_interval(q_left, q_right, left_index)
if q_right > mid: # 在右子树中查询
res_right = self.__query_interval(q_left, q_right, right_index)
return self.function(res_left, res_right) # 返回左右子树元素值的聚合计算结果
# 向上更新实现方法：更新下标为 index 的节点区间值 等于 该节点左右子节点元素值的聚合计算结果
def __pushup(self, index):
left_index = index * 2 + 1 # 左子节点的存储下标
right_index = index * 2 + 2 # 右子节点的存储下标
self.tree[index].val = self.function(self.tree[left_index].val, self.tree[right_index].val)
# 向下更新实现方法：更新下标为 index 的节点所在区间的左右子节点的值和懒惰标记
def __pushdown(self, index):
lazy_tag = self.tree[index].lazy_tag
if lazy_tag is None:
return
left_index = index * 2 + 1 # 左子节点的存储下标
right_index = index * 2 + 2 # 右子节点的存储下标
if self.tree[left_index].lazy_tag is not None:
self.tree[left_index].lazy_tag += lazy_tag # 更新左子节点懒惰标记
else:
self.tree[left_index].lazy_tag = lazy_tag
left_size = (self.tree[left_index].right - self.tree[left_index].left + 1)
self.tree[left_index].val += lazy_tag * left_size # 左子节点每个元素值增加 lazy_tag
if self.tree[right_index].lazy_tag is not None:
self.tree[right_index].lazy_tag += lazy_tag # 更新右子节点懒惰标记
else:
self.tree[right_index].lazy_tag = lazy_tag
right_size = (self.tree[right_index].right - self.tree[right_index].left + 1)
self.tree[right_index].val += lazy_tag * right_size # 右子节点每个元素值增加 lazy_tag
self.tree[index].lazy_tag = None # 更新当前节点的懒惰标记
class Solution:
def getModifiedArray(self, length: int, updates: List[List[int]]) -> List[int]:
nums = [0 for _ in range(length)]
self.ST = SegmentTree(nums, lambda x, y: x + y)
for update in updates:
self.ST.update_interval(update[0], update[1], update[2])
return self.ST.get_nums()