深入理解线段树(Segment Tree)在LeetCode-Go项目中的应用-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00936/article/details/148328043

深入理解线段树(Segment Tree)在LeetCode-Go项目中的应用

LeetCode-Go 该内容是使用Go语言编写的LeetCode题目的完整解决方案集合，实现了100%的测试覆盖率，并且运行时间优于所有题目100%的提交结果。项目地址: https://gitcode.com/gh_mirrors/le/LeetCode-Go

线段树(Segment Tree)是一种非常重要的数据结构，特别适合处理区间查询和更新问题。本文将全面介绍线段树的原理、实现以及在算法问题中的应用。

什么是线段树？

线段树是一种二叉树形数据结构，由Jon Louis Bentley在1977年发明。它主要用于存储区间或线段，并允许快速查询结构内包含某一点的所有区间。

线段树的核心特点：

每个节点代表一个区间
叶子节点代表单个元素区间
内部节点代表其子节点区间的并集
查询时间复杂度为O(log n + k)，其中k是匹配的区间数量

线段树的基本结构

线段树通常使用数组来实现，其结构遵循以下规则：

根节点存储在索引0位置
对于索引i的节点：
- 左子节点索引为2*i+1
- 右子节点索引为2*i+2
叶子节点存储原始数据
内部节点存储合并后的数据

线段树的构建过程采用分治思想，递归地将区间一分为二，直到区间长度为1。

线段树的实现

在LeetCode-Go项目中，线段树的基本实现如下：

type SegmentTree struct {
    data, tree, lazy []int
    left, right      int
    merge            func(i, j int) int
}

func (st *SegmentTree) Init(nums []int, oper func(i, j int) int) {
    // 初始化数据
    st.merge = oper
    data := make([]int, len(nums))
    copy(data, nums)
    st.data = data
    st.tree = make([]int, 4*len(nums))
    st.lazy = make([]int, 4*len(nums))
    if len(nums) > 0 {
        st.buildSegmentTree(0, 0, len(nums)-1)
    }
}

func (st *SegmentTree) buildSegmentTree(treeIndex, left, right int) {
    if left == right {
        st.tree[treeIndex] = st.data[left]
        return
    }
    mid := left + (right-left)/2
    leftIndex := 2*treeIndex + 1
    rightIndex := 2*treeIndex + 2
    st.buildSegmentTree(leftIndex, left, mid)
    st.buildSegmentTree(rightIndex, mid+1, right)
    st.tree[treeIndex] = st.merge(st.tree[leftIndex], st.tree[rightIndex])
}

线段树的查询操作

线段树支持两种查询方式：

1. 直接查询

func (st *SegmentTree) Query(left, right int) int {
    if len(st.data) > 0 {
        return st.queryInTree(0, 0, len(st.data)-1, left, right)
    }
    return 0
}

func (st *SegmentTree) queryInTree(treeIndex, l, r, queryL, queryR int) int {
    if l == queryL && r == queryR {
        return st.tree[treeIndex]
    }
    mid := l + (r-l)/2
    leftIndex := 2*treeIndex + 1
    rightIndex := 2*treeIndex + 2
    if queryL > mid {
        return st.queryInTree(rightIndex, mid+1, r, queryL, queryR)
    } else if queryR <= mid {
        return st.queryInTree(leftIndex, l, mid, queryL, queryR)
    }
    return st.merge(
        st.queryInTree(leftIndex, l, mid, queryL, mid),
        st.queryInTree(rightIndex, mid+1, r, mid+1, queryR),
    )
}

2. 懒查询（配合懒更新）

func (st *SegmentTree) QueryLazy(left, right int) int {
    if len(st.data) > 0 {
        return st.queryLazyInTree(0, 0, len(st.data)-1, left, right)
    }
    return 0
}

func (st *SegmentTree) queryLazyInTree(treeIndex, l, r, queryL, queryR int) int {
    mid := l + (r-l)/2
    leftIndex := 2*treeIndex + 1
    rightIndex := 2*treeIndex + 2
    
    // 处理懒标记
    if st.lazy[treeIndex] != 0 {
        st.tree[treeIndex] += (r-l+1)*st.lazy[treeIndex]
        if l != r {
            st.lazy[leftIndex] += st.lazy[treeIndex]
            st.lazy[rightIndex] += st.lazy[treeIndex]
        }
        st.lazy[treeIndex] = 0
    }
    
    // 查询逻辑
    if queryL <= l && queryR >= r {
        return st.tree[treeIndex]
    }
    if queryL > mid {
        return st.queryLazyInTree(rightIndex, mid+1, r, queryL, queryR)
    }
    if queryR <= mid {
        return st.queryLazyInTree(leftIndex, l, mid, queryL, queryR)
    }
    return st.merge(
        st.queryLazyInTree(leftIndex, l, mid, queryL, mid),
        st.queryLazyInTree(rightIndex, mid+1, r, mid+1, queryR),
    )
}

线段树的更新操作

1. 单点更新

func (st *SegmentTree) Update(index, val int) {
    if len(st.data) > 0 {
        st.updateInTree(0, 0, len(st.data)-1, index, val)
    }
}

func (st *SegmentTree) updateInTree(treeIndex, l, r, index, val int) {
    if l == r {
        st.tree[treeIndex] = val
        return
    }
    mid := l + (r-l)/2
    leftIndex := 2*treeIndex + 1
    rightIndex := 2*treeIndex + 2
    if index > mid {
        st.updateInTree(rightIndex, mid+1, r, index, val)
    } else {
        st.updateInTree(leftIndex, l, mid, index, val)
    }
    st.tree[treeIndex] = st.merge(st.tree[leftIndex], st.tree[rightIndex])
}

2. 区间更新（使用懒标记）

func (st *SegmentTree) UpdateLazy(updateL, updateR, val int) {
    if len(st.data) > 0 {
        st.updateLazyInTree(0, 0, len(st.data)-1, updateL, updateR, val)
    }
}

func (st *SegmentTree) updateLazyInTree(treeIndex, l, r, updateL, updateR, val int) {
    mid := l + (r-l)/2
    leftIndex := 2*treeIndex + 1
    rightIndex := 2*treeIndex + 2
    
    // 先处理懒标记
    if st.lazy[treeIndex] != 0 {
        st.tree[treeIndex] += (r-l+1)*st.lazy[treeIndex]
        if l != r {
            st.lazy[leftIndex] += st.lazy[treeIndex]
            st.lazy[rightIndex] += st.lazy[treeIndex]
        }
        st.lazy[treeIndex] = 0
    }
    
    // 当前区间不在更新范围内
    if l > r || l > updateR || r < updateL {
        return
    }
    
    // 当前区间完全在更新范围内
    if updateL <= l && updateR >= r {
        st.tree[treeIndex] += (r-l+1)*val
        if l != r {
            st.lazy[leftIndex] += val
            st.lazy[rightIndex] += val
        }
        return
    }
    
    // 更新子区间
    st.updateLazyInTree(leftIndex, l, mid, updateL, updateR, val)
    st.updateLazyInTree(rightIndex, mid+1, r, updateL, updateR, val)
    st.tree[treeIndex] = st.merge(st.tree[leftIndex], st.tree[rightIndex])
}