Python数据结构---线段树

麻辣小兔喵

于 2023-03-17 23:28:50 发布

阅读量485

点赞数

分类专栏： Python 文章标签： python 数据结构算法

本文链接：https://blog.csdn.net/weixin_51705589/article/details/129630169

版权

Python 专栏收录该内容

17 篇文章 2 订阅

订阅专栏

线段树（Segment Tree）是一种常见的数据结构，常用于解决区间查询问题，例如区间最大值、最小值、区间和、区间平均值等。

下面讲一下如何在Python 中实现线段树。

1.定义节点类

首先，我们需要定义一个节点类来表示线段树中的节点。每个节点包含一个左右子节点、区间起始和结束位置以及一些其他信息（如区间和、区间最大值、最小值等）。

在 Python 中，可以使用以下代码定义一个节点类：

class SegmentTreeNode:
    def __init__(self, start, end):
        self.start = start  # 区间起始位置
        self.end = end  # 区间结束位置
        self.left = None  # 左子节点
        self.right = None  # 右子节点
        self.sum = 0  # 区间和
        # 其他信息（例如区间最大值、最小值等）可以根据需要添加

2.构建线段树

构建线段树的过程可以使用递归的方式实现。
首先，定义一个 build 函数，该函数接收一个数组以及区间起始和结束位置，返回构建好的线段树的根节点。
在函数内部，首先创建一个新节点表示当前区间，然后递归构建左右子树，最后将左右子树的区间和相加作为当前节点的区间和。

def build(nums, start, end):
    if start > end:
        return None
    root = SegmentTreeNode(start, end)
    if start == end:
        root.sum = nums[start]
        return root
    mid = (start + end) // 2
    root.left = build(nums, start, mid)
    root.right = build(nums, mid+1, end)
    root.sum = root.left.sum + root.right.sum
    return root

3.区间查询

线段树最常见的用途就是区间查询。
例如，查询一个区间的和、最大值、最小值等。
以下是一个实现区间查询的示例代码：

def query(root, start, end):
    if not root:
        return 0
    if start <= root.start and end >= root.end:
        return root.sum
    mid = (root.start + root.end) // 2
    res = 0
    if start <= mid:
        res += query(root.left, start, end)
    if end > mid:
        res += query(root.right, start, end)
    return res

这个函数接收一个根节点、区间起始和结束位置，返回区间和。如果当前节点的区间包含在查询区间内，则直接返回当前节点的区间和；否则，将查询区间分别递归到左右子树中查询，并将左右子树的查询结果相加作为最终结果。

4.区间修改

除了区间查询之外，线段树还可以用来实现区间修改。例如，修改一个区间中的数值，或者给一个区间加上一个定值。
以下是一个实现区间修改的示例代码：

def update(root, start, end, val):
    if not root:
        return
    if start <= root.start and end >= root.end:
        root.sum += val
        return
    mid = (root.start + root.end) // 2
    if start <= mid:
        update(root.left, start, end, val)
    if end > mid:
        update(root.right, start, end, val)
    root.sum = root.left.sum + root.right.sum