这篇文章是向 Ray Wenderlich 中 Swift Algorithm Club 的投稿文。用来讲述
Segment Tree的区间更新操作以及 Swift 实现方案。Open Issue #591
在原始的实现中，Artur Antonov 利用 Swift 中很优雅的泛型<T>方式来对 Segment Tree 这个数据结构进行了实现。并且在这个实现中可以对其传递一个 function: (T, T) -> T 来反映两个叶子节点对父节点的影响。并且泛型的优势体现在可以用 Segment Tree 来解决字符串的问题，就像 Playground 中最后的这个例子：
stringSegmentTree.replaceItem(at: 0, withItem: "I")
stringSegmentTree.replaceItem(at: 1, withItem: " like")
stringSegmentTree.replaceItem(at: 2, withItem: " algorithms")
stringSegmentTree.replaceItem(at: 3, withItem: " and")
stringSegmentTree.replaceItem(at: 4, withItem: " swift")
stringSegmentTree.replaceItem(at: 5, withItem: "!")
print(stringSegmentTree.query(leftBound: 0, rightBound: 5))
// "I like algorithms and swift!"<T> 这种方式固然有趣，可以利用 Segment Tree 来做字符串处理的地方真的少之又少。况且此数据结构的出现源自于区间性处理问题，将一维集合通过分制思想多线段划分，从而以空间换时间达到高效操作的数据结构。而最常见的问题即为解决区间内任意子区间求和的问题。从原本 O(n) 时间复杂度以二分查询的形式降低为 O(logn)，达到高效查询。
pushUp 向下递归，向上更新操作
我们可以参考一下 Artur Antonov 的实现方法，其中有建树、单点更新和区间查询三种操作。这里将 Segment Tree 的建树操作和单点更新拿出来看一下：
public init(array: [T], leftBound: Int, rightBound: Int, function: @escaping (T, T) -> T) {
self.leftBound = leftBound
self.rightBound = rightBound
self.function = function
// ①
if leftBound == rightBound {
value = array[leftBound]
}
// ②
else {
let middle = (leftBound + rightBound) / 2
leftChild = SegmentTree<T>(array: array, leftBound: leftBound, rightBound: middle, function: function)
rightChild = SegmentTree<T>(array: array, leftBound: middle+1, rightBound: rightBound, function: function)
value = function(leftChild!.value, rightChild!.value)
}
}在 ① 位置中，当遍历节点的左边极限范围和右边极限范围相同时，则说明这是一个叶子节点，则直接为其赋值即可。而 ② 位置说明该节点为父节点（仅对于有无子节点而言）时，需要对下层叶子继续构建，在后续的递归过程时对该节点进行构建。
继而，我们来看单点更新(Query)的操作：
public func query(leftBound: Int, rightBound: Int) -> T {
if self.leftBound == leftBound && self.rightBound == rightBound {
return self.value
}
guard let leftChild = leftChild else { fatalError("leftChild should not be nil") }
guard let rightChild = rightChild else { fatalError("rightChild should not be nil") }
// ①
if leftChild.rightBound < leftBound {
return rightChild.query(leftBound: leftBound, rightBound: rightBound)
}
// ②
else if rightChild.leftBound > rightBound {
return leftChild.query(leftBound: leftBound, rightBound: rightBound)
}
// ③
else {
let leftResult = leftChild.query(leftBound: leftBound, rightBound: leftChild.rightBound)
let rightResult = rightChild.query(leftBound:rightChild.leftBound, rightBound: rightBound)
return function(leftResult, rightResult)
}
}在 ① 位置中是查询位置的左边界在当前节点的右边界，这时候说明击中范围偏右，需要继续向右递归查询。② 位置的情况和 ① 的情况相反，继续想左递归查询即可。而 ③ 中的代码说明当前位置为需查询位置，则继续向下递归查询，直到为单值或是恰好卡中区间即可返回结果。
从上述两个代码中，我们发现这么一个共有的部分，就是向下递归，向上更新部分。其实，我们完全可以将这个向下更新的操作分离处理，当需要的时候调用即可。以下是 func pushUp(lson: LazySegmentTree, rson: LazySegmentTree) 操作：
// MARK: - Push Up Operation
// Description: 这里是 push up 操作，用来对 Segment Tree 向上更新
private func pushUp(lson: LazySegmentTree, rson: LazySegmentTree) {
self.value = lson.value + rson.value
}以上代码中，仅仅对 SumSegmentTree 进行了实现，即求和线段树。有了这个方法，我们队构建和单点替换操作就可以直接调用实现。示例代码如下：
public init(array: [Int], leftBound: Int, rightBound: Int) {
self.leftBound = leftBound
self.rightBound = rightBound
self.value = 0
self.lazyValue = 0
if leftBound == rightBound {
value = array[leftBound]
return
}
let middle = (leftBound + rightBound) / 2
leftChild = LazySegmentTree(array: array, leftBound: leftBound, rightBound: middle)
rightChild = LazySegmentTree(array: array, leftBound: middle + 1, rightBound: rightBound)
pushUp(lson: leftChild!, rson: rightChild!)
}
// MARK: - One Item Update
public func update(index: Int, newValue: Int) {
if self.leftBound == self.rightBound {
self.value = newValue
return
}
guard let leftChild = leftChild else { fatalError("leftChild should not be nil") }
guard let rightChild = rightChild else { fatalError("rightChild should not be nil") }
let middle = (self.leftBound + self.rightBound) / 2
if index <= middle { leftChild.update(index: index, newValue: newValue) }
else { rightChild.update(index: index, newValue: newValue) }
pushUp(lson: leftChild, rson: rightChild)
}从 pushDown 操作到使用 Lazy 思想对区间进行更新
上文中对于 pushUp 操作也许你会觉得其实没有多少变化。其实是为了引入接下来要说的 pushDown，进行相对而言的。
但是我想先来聊聊区间操作。所谓区间操作就是对集合中连续的一系列数据进行更新操作。在 Artur Antonov 的代码中，并没有实现这个操作，但是我们可以直接想到，通过一个 for 循环就可完成对于一个区间进行更新：
// Sample: 对集合下标 [2, 5] 的数据进行更新
for index in 2 ... 5 {
sumSegmentTree.update(index: index, incremental: 3)
}可是这要就会存在一个 O(n) 复杂度的遍历，最后使得区间更新操作的时间复杂度为 O(nlogn)。那么有没有一种方法，将其优化至 O(logn)，从而使其和单点更新以及范围查询两个操作一样优美呢？
考虑这个问题，我们需要重新来审视线段树的这个结构：
而在代码中，我们只能拿到这棵线段树的根节点。如果我们想深入到底层，然后再对树从最底层的叶子节点进行 pushUp 向上更新这样就能完成当前目标了，这样就要求先对这棵二叉树进行一次遍历。但一次遍历可不是易事，无论是先序还是后序，无论 DFS 还是 BFS，找到每个叶子节点一定是需要 O(n) 的复杂度，这已经超出了我们的期望。
所以我们开始考虑 pushDown 向下更新这种新操作，我们在大范围的父内对数据进行更新，然后再向下递归时，根据更新要求，再对叶子节点进行更新就好了。但是，这好像与 pushUp 没有什么改善，整棵树更新完也需要 O(n) 的复杂度。继续思考，当我们做范围查询操作时，也是从根节点开始，向下去卡区间范围进行递归查询。既然查询和更新的方向是一致的，我们就可以带着更新来查询了不是吗？😁 举个例子：
update 操作是为了让下标为 1 到 3 这三个元素的值加 2，所以我们在 2 层的第一个节点和 3 层第三个节点中加入 lazy 标记，代表这些节点需要更新，并在这个过程中，由于我们无需在 root 节点增加标记，所以在递归之前执行一次 pushDown 操作来更新 root 节点的值。
query 操作时为了查询第一个元素，我们逐层递归遍历，当来到 2 层第一个节点的时候发现，这个节点 lazy 标记不为空，所以更新一下它的值，来到 3 层第一个节点时也相同，更新之后查出结果。
是的，这就是 Lazy 思想。当需要更新时，我们只对指定范围的父亲节点中增加一个待更新标记，说明这些范围我们需要更新。而在查询的时候，如果发现这些标记，则再对其子节点进行更新。而这个向下更新操作就是 pushDown 方法。
下面我们用 Swift 来实现一下这个方案，其案例是对于 SumSegmentTree 的完整代码：
public class LazySegmentTree {
private var value: Int
private var leftBound: Int
private var rightBound: Int
private var leftChild: LazySegmentTree?
private var rightChild: LazySegmentTree?
// Interval Update Lazy Element
private var lazyValue: Int
// MARK: - Push Up Operation
// Description: 这里是 push up 操作，用来对 Segment Tree 向上更新
private func pushUp(lson: LazySegmentTree, rson: LazySegmentTree) {
self.value = lson.value + rson.value
}
// MARK: - Push Down Operation
// Description: 这里是 push down 操作，用来对 Segment Tree 向下更新
// Open Interface Function: 此处应该开放方法对齐进行 Override
private func pushDown(round: Int, lson: LazySegmentTree, rson: LazySegmentTree) {
if lazyValue != 0 {
lson.lazyValue += lazyValue
rson.lazyValue += lazyValue
lson.value += lazyValue * (round - (round >> 1))
rson.value += lazyValue * (round >> 1)
lazyValue = 0
}
}
public init(array: [Int], leftBound: Int, rightBound: Int) {
self.leftBound = leftBound
self.rightBound = rightBound
self.value = 0
self.lazyValue = 0
if leftBound == rightBound {
value = array[leftBound]
return
}
let middle = (leftBound + rightBound) / 2
leftChild = LazySegmentTree(array: array, leftBound: leftBound, rightBound: middle)
rightChild = LazySegmentTree(array: array, leftBound: middle + 1, rightBound: rightBound)
pushUp(lson: leftChild!, rson: rightChild!)
}
public convenience init(array: [Int]) {
self.init(array: array, leftBound: 0, rightBound: array.count - 1)
}
public func query(leftBound: Int, rightBound: Int) -> Int {
if leftBound <= self.leftBound && self.rightBound <= rightBound {
return value
}
guard let leftChild = leftChild else { fatalError("leftChild should not be nil") }
guard let rightChild = rightChild else { fatalError("rightChild should not be nil") }
pushDown(round: self.rightBound - self.leftBound + 1, lson: leftChild, rson: rightChild)
let middle = (self.leftBound + self.rightBound) / 2
var result: Int = 0
if leftBound <= middle { result += leftChild.query(leftBound: leftBound, rightBound: rightBound) }
if rightBound > middle { result += rightChild.query(leftBound: leftBound, rightBound: rightBound) }
return result
}
// MARK: - One Item Update
public func update(index: Int, incremental: Int) {
if self.leftBound == self.rightBound {
self.value += incremental
return
}
guard let leftChild = leftChild else { fatalError("leftChild should not be nil") }
guard let rightChild = rightChild else { fatalError("rightChild should not be nil") }
let middle = (self.leftBound + self.rightBound) / 2
if index <= middle { leftChild.update(index: index, incremental: incremental) }
else { rightChild.update(index: index, incremental: incremental) }
pushUp(lson: leftChild, rson: rightChild)
}
// MARK: - Interval Item Update
public func update(leftBound: Int, rightBound: Int, incremental: Int) {
if leftBound <= self.leftBound && self.rightBound <= rightBound {
self.lazyValue += incremental
self.value += incremental * (self.rightBound - self.leftBound + 1)
return
}
guard let leftChild = leftChild else { fatalError() }
guard let rightChild = rightChild else { fatalError() }
pushDown(round: self.rightBound - self.leftBound + 1, lson: leftChild, rson: rightChild)
let middle = (self.leftBound + self.rightBound) / 2
if leftBound <= middle { leftChild.update(leftBound: leftBound, rightBound: rightBound, incremental: incremental) }
if middle < rightBound { rightChild.update(leftBound: leftBound, rightBound: rightBound, incremental: incremental) }
pushUp(lson: leftChild, rson: rightChild)
}
}我们来解释一下这些代码片段：
private var lazyValue: Int这里我们增加了一个 lazyValue 变量，这个变量的含义在当前数据结构中是当前节点需要的增量。因为其 update 操作时用来对指定元素和指定范围的所有元素进行数量的增减。lazyValue 的值弱不为0，则说明当前节点的 value 需要更新，而其真值为 value + lazyValue * (rightBound - leftBound + 1)。
// MARK: - Push Down Operation
// Description: 这里是 push down 操作，用来对 Segment Tree 向下更新
// Open Interface Function: 此处应该开放方法对齐进行 Override
private func pushDown(round: Int, lson: LazySegmentTree, rson: LazySegmentTree) {
if lazyValue != 0 {
lson.lazyValue += lazyValue
rson.lazyValue += lazyValue
lson.value += lazyValue * (round - (round >> 1))
rson.value += lazyValue * (round >> 1)
lazyValue = 0
}
}
这块代码主要是检查当前节点是否需要更新，如果 lazyValue 的值不为 0，则需要向下更新一次。其更新原则就是左右范围差值算出元素个数，乘以增量算出父亲节点的增量，之后更新子节点的增量值，最后清零自身增量值。
public func query(leftBound: Int, rightBound: Int) -> Int {
if leftBound <= self.leftBound && self.rightBound <= rightBound {
return value
}
guard let leftChild = leftChild else { fatalError("leftChild should not be nil") }
guard let rightChild = rightChild else { fatalError("rightChild should not be nil") }
pushDown(round: self.rightBound - self.leftBound + 1, lson: leftChild, rson: rightChild)
let middle = (self.leftBound + self.rightBound) / 2
var result: Int = 0
if leftBound <= middle { result += leftChild.query(leftBound: leftBound, rightBound: rightBound) }
if rightBound > middle { result += rightChild.query(leftBound: leftBound, rightBound: rightBound) }
return result
}根据前文的分析，当对区间进行查询的时候，我们就调用一次 pushDown 来向下更新数据，以保证在递归查询之前所有的数据处于最新状态。有的读者可能有疑问，为什么这里的代码比原版的少了一个 if else 的判断？因为我使用的是判断中值与边界的关系，而不是边界与边界的关系。笔者认为这么判断代码会更加简洁，思路也会更加清晰。
以上便是 lazy 思想实现线段树的区间更新操作及思路。作为线段树的扩展，还可以通过可持久化线段树（也称主席树）对于永久化保存状态的概念，更深一步的理解 lazy 思想。另外，在代码实现的时候，notonlysuccess大牛的实例代码，可以说是线段树中最为清晰的实现方式，给出链接有兴趣可以自行学习。
后记
线段树在工程项目中是一科极少出现的数据结构，由于它属于区间树（Interval Tree）的（而且是一种大开销的）一种实现，所以更多的用来解决数学中的区间问题和统计问题中。在之后的完善中，还会针对于 SumSegmentTree 增加更多的区间（交、并、补）等操作，以及叶子节点合并、二维线段树等等更为复杂，希望大家共同学习，一起完善 Segment Tree In Swift 🐦。
文章最后求一个中译英的译者，一起向
swift-algorithm-club进行 contribution。
