线段树如何实现区间加和区间查询非零数个数？

Question

线段树是否能维护区间加，和区间查询不为零的数的个数？为什么？如果是，那么该如何实现？

已经向 百度 AI 和 gpt_4o_free 询问过，得到肯定结果，但我仍然不明白为什么。

Answer 1

线段树能否维护区间加和区间查询非零数个数

线段树是可以维护区间加和区间查询非零数个数的。原因在于线段树是一种二叉树数据结构，它可以将一个区间划分为多个小子区间，并且通过树的节点来存储和管理这些子区间的信息。对于区间加操作，我们可以使用懒标记（Lazy Propagation）的技术，避免每次操作都对整个区间的所有元素进行更新，而是将更新信息暂时存储在节点上，在需要查询或进一步更新时再将标记下推到子节点。对于区间查询非零数个数，我们可以在每个节点上维护该子区间内非零数的数量，通过合并子节点的信息来得到父节点的信息。

实现思路

1. 线段树节点定义

每个节点需要存储以下信息：

• 区间范围：记录该节点所代表的区间的左右端点。
• 非零数个数：该节点所代表的区间内非零数的数量。
• 懒标记：用于区间加操作，记录该节点需要向下传递的增量。

class SegmentTreeNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
        self.count_non_zero = 0
        self.lazy = 0
        self.left_child = None
        self.right_child = None

2. 构建线段树

根据给定的数组构建线段树，递归地将区间划分为子区间，并计算每个节点的非零数个数。

def build_tree(arr, left, right):
    node = SegmentTreeNode(left, right)
    if left == right:
        node.count_non_zero = 1 if arr[left] != 0 else 0
        return node
    mid = (left + right) // 2
    node.left_child = build_tree(arr, left, mid)
    node.right_child = build_tree(arr, mid + 1, right)
    node.count_non_zero = node.left_child.count_non_zero + node.right_child.count_non_zero
    return node

3. 下推懒标记

当需要访问某个节点时，如果该节点有懒标记，需要将标记下推到子节点，并更新子节点的信息。

def push_down(node):
    if node.lazy != 0:
        if node.left_child:
            node.left_child.lazy += node.lazy
            if node.left_child.lazy != 0:
                node.left_child.count_non_zero = node.left_child.right - node.left_child.left + 1
            else:
                # 重新计算非零数个数，这里需要根据具体情况实现
                pass
        if node.right_child:
            node.right_child.lazy += node.lazy
            if node.right_child.lazy != 0:
                node.right_child.count_non_zero = node.right_child.right - node.right_child.left + 1
            else:
                # 重新计算非零数个数，这里需要根据具体情况实现
                pass
        node.lazy = 0

4. 区间加操作

对指定区间内的所有元素加上一个值，使用懒标记来优化更新操作。

def update_range(node, left, right, val):
    if left > node.right or right < node.left:
        return
    if left <= node.left and right >= node.right:
        node.lazy += val
        if node.lazy != 0:
            node.count_non_zero = node.right - node.left + 1
        else:
            # 重新计算非零数个数，这里需要根据具体情况实现
            pass
        return
    push_down(node)
    mid = (node.left + node.right) // 2
    if left <= mid:
        update_range(node.left_child, left, right, val)
    if right > mid:
        update_range(node.right_child, left, right, val)
    node.count_non_zero = node.left_child.count_non_zero + node.right_child.count_non_zero

5. 区间查询非零数个数

查询指定区间内非零数的个数，递归地合并子节点的信息。

def query_range(node, left, right):
    if left > node.right or right < node.left:
        return 0
    if left <= node.left and right >= node.right:
        return node.count_non_zero
    push_down(node)
    mid = (node.left + node.right) // 2
    result = 0
    if left <= mid:
        result += query_range(node.left_child, left, right)
    if right > mid:
        result += query_range(node.right_child, left, right)
    return result

示例代码调用

# 示例数组
arr = [1, 0, 2, 0, 3]
root = build_tree(arr, 0, len(arr) - 1)

# 区间加操作
update_range(root, 1, 3, 1)

# 区间查询非零数个数
count = query_range(root, 0, len(arr) - 1)
print("区间内非零数的个数:", count)

复杂度分析

• 时间复杂度：构建线段树的时间复杂度为 $O(n)$，其中 $n$ 是数组的长度。区间加和区间查询操作的时间复杂度均为 $O(\log n)$。
• 空间复杂度：线段树的空间复杂度为 $O(n)$。

通过上述实现，我们可以使用线段树有效地维护区间加和区间查询非零数个数的操作。需要注意的是，在实际应用中，对于下推懒标记时重新计算非零数个数的部分，需要根据具体情况进行更详细的实现，以处理更复杂的情况。