带区间更新的数组零元素统计问题高效解法问询

问题背景 #

给定初始全为0的数组（长度n<10^{5），需执行q次操作（q<2*10}5）：

• 区间加法：给[l,r]内所有元素加x（x可正可负）
• 零元素查询：统计[l,r]内值为0的元素个数

暴力遍历的O(qn)复杂度无法满足性能要求，需实现O(qlog(n))的高效解法。

线段树节点设计 #

每个节点维护4个核心字段：

• zero_cnt：当前区间内值为0的元素数量
• lazy：待向下传递的区间增量（懒标记，避免重复递归更新子节点）
• min_val：当前区间元素的最小值
• max_val：当前区间元素的最大值

字段作用：

• min_val/max_val用于剪枝：若区间min_val>0或max_val<0，直接返回0；若两者都为0，直接返回区间长度，无需递归到子节点。
• 懒标记将区间操作的影响延迟到需要访问子节点时再传递，保证操作的对数时间复杂度。

核心操作实现 #

1. 懒标记传递（push_down） #

当需要访问当前节点的子节点时，先将lazy标记传递给子节点，更新子节点的状态：

def push_down(node, node_l, node_r):
    if node.lazy == 0:
        return
    val = node.lazy
    mid = (node_l + node_r) // 2
    left = node.left
    right = node.right

    # 更新左子节点
    left.min_val += val
    left.max_val += val
    left.lazy += val
    # 若左子节点所有元素值相同，直接更新zero_cnt
    if left.min_val == left.max_val:
        left.zero_cnt = mid - node_l + 1 if left.min_val == 0 else 0

    # 更新右子节点
    right.min_val += val
    right.max_val += val
    right.lazy += val
    if right.min_val == right.max_val:
        right.zero_cnt = node_r - mid if right.min_val == 0 else 0

    # 清空当前节点的懒标记
    node.lazy = 0

2. 区间加法操作（update） #

递归更新目标区间的状态，利用懒标记避免不必要的递归：

def update(node, node_l, node_r, target_l, target_r, x):
    if target_r < node_l or target_l > node_r:
        return
    # 当前区间完全被目标区间覆盖，直接更新
    if target_l <= node_l and node_r <= target_r:
        node.min_val += x
        node.max_val += x
        node.lazy += x
        # 若当前区间所有元素值相同，更新zero_cnt
        if node.min_val == node.max_val:
            node.zero_cnt = node_r - node_l + 1 if node.min_val == 0 else 0
        return
    # 先传递懒标记，再递归更新子节点
    push_down(node, node_l, node_r)
    mid = (node_l + node_r) // 2
    update(left, node_l, mid, target_l, target_r, x)
    update(right, mid+1, node_r, target_l, target_r, x)
    # 合并子节点的zero_cnt
    node.zero_cnt = left.zero_cnt + right.zero_cnt

3. 零元素查询操作（query） #

递归统计目标区间内的零元素个数，同样利用剪枝优化：

def query(node, node_l, node_r, target_l, target_r):
    if target_r < node_l or target_l > node_r:
        return 0
    # 当前区间完全被目标区间覆盖，直接返回统计值
    if target_l <= node_l and node_r <= target_r:
        return node.zero_cnt
    # 传递懒标记后，递归查询子节点
    push_down(node, node_l, node_r)
    mid = (node_l + node_r) // 2
    left_cnt = query(left, node_l, mid, target_l, target_r)
    right_cnt = query(right, mid+1, node_r, target_l, target_r)
    return left_cnt + right_cnt

示例运行流程 #

初始数组：[0, 0, 0, 0]（索引0-3）

1. 执行区间加法：给[2,3]加3
   • 对应区间的min/max变为3，zero_cnt置为0；数组变为[0,0,3,3]
2. 查询[1,2]的零元素个数：
   • 递归查询后，索引1的zero_cnt为1，索引2的zero_cnt为0，总和为1
3. 执行区间加法：给[3,3]加-3
   • 对应叶子节点的min/max变为0，zero_cnt置为1；数组变为[0,0,3,0]
4. 查询[0,3]的零元素个数：
   • 索引0、1、3的zero_cnt均为1，总和为3

注意事项 #

• 懒标记必须在访问子节点前传递，确保子节点状态最新
• 仅当区间内所有元素值相同时，才能直接计算zero_cnt；否则需依赖子节点的统计结果合并
• 线段树初始化时，叶子节点的min/max为0，zero_cnt为1；非叶子节点的min/max为0，zero_cnt为左右子节点之和，lazy初始为0

内容的提问来源于stack exchange，提问作者dispenomi