线段树——区间查询优化利器

后端

2023-11-07 18:33:51

线段树概述

线段树是一种树形数据结构，每个节点代表一个数组中的某个区间。线段树的每个节点都有三个子节点：左子节点、右子节点和父节点。左子节点和右子节点分别代表区间的前半部分和后半部分，父节点代表整个区间。

线段树的构建过程如下：

将数组中的每个元素作为叶节点，从左到右依次插入到线段树中。
将叶节点的父节点设为叶节点的左子节点和右子节点的父节点。
重复步骤2，直到所有的节点都具有父节点。

线段树的查询过程如下：

从根节点开始，依次比较查询区间与当前节点所代表的区间。
如果查询区间完全包含在当前节点所代表的区间内，则直接返回当前节点的值。
如果查询区间与当前节点所代表的区间有重叠，则递归查询当前节点的左子节点和右子节点。
将左子节点和右子节点的查询结果合并，得到最终的查询结果。

线段树的修改过程如下：

从根节点开始，依次比较修改区间与当前节点所代表的区间。
如果修改区间完全包含在当前节点所代表的区间内，则直接修改当前节点的值。
如果修改区间与当前节点所代表的区间有重叠，则递归修改当前节点的左子节点和右子节点。
将左子节点和右子节点的修改结果合并，得到最终的修改结果。

线段树的应用

线段树可以用来解决许多与区间相关的复杂问题，例如：

区间和：计算一个数组中某个区间内所有元素的和。
区间最大值：计算一个数组中某个区间内所有元素的最大值。
区间最小值：计算一个数组中某个区间内所有元素的最小值。
区间中位数：计算一个数组中某个区间内所有元素的中位数。
区间众数：计算一个数组中某个区间内出现次数最多的元素。
区间乘积：计算一个数组中某个区间内所有元素的乘积。
区间异或和：计算一个数组中某个区间内所有元素的异或和。

Python实现

class Node:
    def __init__(self, start, end, value):
        self.start = start
        self.end = end
        self.value = value
        self.left = None
        self.right = None

class SegmentTree:
    def __init__(self, array):
        self.root = self._build_tree(array, 0, len(array) - 1)

    def _build_tree(self, array, start, end):
        if start == end:
            return Node(start, end, array[start])

        mid = (start + end) // 2
        left_child = self._build_tree(array, start, mid)
        right_child = self._build_tree(array, mid + 1, end)

        return Node(start, end, left_child.value + right_child.value)

    def query(self, start, end):
        return self._query(self.root, start, end)

    def _query(self, node, start, end):
        if node.start == start and node.end == end:
            return node.value

        mid = (node.start + node.end) // 2

        if end <= mid:
            return self._query(node.left, start, end)
        elif start > mid:
            return self._query(node.right, start, end)
        else:
            return self._query(node.left, start, mid) + self._query(node.right, mid + 1, end)

    def update(self, index, value):
        self._update(self.root, index, value)

    def _update(self, node, index, value):
        if node.start == node.end == index:
            node.value = value
            return

        mid = (node.start + node.end) // 2

        if index <= mid:
            self._update(node.left, index, value)
        else:
            self._update(node.right, index, value)

        node.value = node.left.value + node.right.value

if __name__ == "__main__":
    array = [1, 2, 3, 4, 5]
    segment_tree = SegmentTree(array)

    print(segment_tree.query(0, 2))  # 6
    print(segment_tree.query(1, 3))  # 9
    print(segment_tree.query(2, 4))  # 12

    segment_tree.update(2, 10)

    print(segment_tree.query(0, 2))  # 11
    print(segment_tree.query(1, 3))  # 15
    print(segment_tree.query(2, 4))  # 17