1+ public class LazySegmentTree {
2+
3+ private var value : Int
4+
5+ private var leftBound : Int
6+
7+ private var rightBound : Int
8+
9+ private var leftChild : LazySegmentTree ?
10+
11+ private var rightChild : LazySegmentTree ?
12+
13+ // Interval Update Lazy Element
14+ private var lazyValue : Int
15+
16+ // MARK: - Push Up Operation
17+ // Description: pushUp() - update items to the top
18+ private func pushUp( lson: LazySegmentTree , rson: LazySegmentTree ) {
19+ self . value = lson. value + rson. value
20+ }
21+
22+ // MARK: - Push Down Operation
23+ // Description: pushDown() - update items to the bottom
24+ private func pushDown( round: Int , lson: LazySegmentTree , rson: LazySegmentTree ) {
25+ guard lazyValue != 0 else { return }
26+ lson. lazyValue += lazyValue
27+ rson. lazyValue += lazyValue
28+ lson. value += lazyValue * ( round - ( round >> 1 ) )
29+ rson. value += lazyValue * ( round >> 1 )
30+ lazyValue = 0
31+ }
32+
33+ public init ( array: [ Int ] , leftBound: Int , rightBound: Int ) {
34+ self . leftBound = leftBound
35+ self . rightBound = rightBound
36+ self . value = 0
37+ self . lazyValue = 0
38+
39+ guard leftBound != rightBound else {
40+ value = array [ leftBound]
41+ return
42+ }
43+
44+ let middle = leftBound + ( rightBound - leftBound) / 2
45+ leftChild = LazySegmentTree ( array: array, leftBound: leftBound, rightBound: middle)
46+ rightChild = LazySegmentTree ( array: array, leftBound: middle + 1 , rightBound: rightBound)
47+ if let leftChild = leftChild, let rightChild = rightChild {
48+ pushUp ( lson: leftChild, rson: rightChild)
49+ }
50+ }
51+
52+ public convenience init ( array: [ Int ] ) {
53+ self . init ( array: array, leftBound: 0 , rightBound: array. count - 1 )
54+ }
55+
56+ public func query( leftBound: Int , rightBound: Int ) -> Int {
57+ if leftBound <= self . leftBound && self . rightBound <= rightBound {
58+ return value
59+ }
60+ guard let leftChild = leftChild else { fatalError ( " leftChild should not be nil " ) }
61+ guard let rightChild = rightChild else { fatalError ( " rightChild should not be nil " ) }
62+
63+ pushDown ( round: self . rightBound - self . leftBound + 1 , lson: leftChild, rson: rightChild)
64+
65+ let middle = self . leftBound + ( self . rightBound - self . leftBound) / 2
66+ var result = 0
67+
68+ if leftBound <= middle { result += leftChild. query ( leftBound: leftBound, rightBound: rightBound) }
69+ if rightBound > middle { result += rightChild. query ( leftBound: leftBound, rightBound: rightBound) }
70+
71+ return result
72+ }
73+
74+ // MARK: - One Item Update
75+ public func update( index: Int , incremental: Int ) {
76+ guard self . leftBound != self . rightBound else {
77+ self . value += incremental
78+ return
79+ }
80+ guard let leftChild = leftChild else { fatalError ( " leftChild should not be nil " ) }
81+ guard let rightChild = rightChild else { fatalEr
F438
ror ( " rightChild should not be nil " ) }
82+
83+ let middle = self . leftBound + ( self . rightBound - self . leftBound) / 2
84+
85+ if index <= middle { leftChild. update ( index: index, incremental: incremental) }
86+ else { rightChild. update ( index: index, incremental: incremental) }
87+ pushUp ( lson: leftChild, rson: rightChild)
88+ }
89+
90+ // MARK: - Interval Item Update
91+ public func update( leftBound: Int , rightBound: Int , incremental: Int ) {
92+ if leftBound <= self . leftBound && self . rightBound <= rightBound {
93+ self . lazyValue += incremental
94+ self . value += incremental * ( self . rightBound - self . leftBound + 1 )
95+ return
96+ }
97+
98+ guard let leftChild = leftChild else { fatalError ( " leftChild should not be nil " ) }
99+ guard let rightChild = rightChild else { fatalError ( " rightChild should not be nil " ) }
100+
101+ pushDown ( round: self . rightBound - self . leftBound + 1 , lson: leftChild, rson: rightChild)
102+
103+ let middle = self . leftBound + ( self . rightBound - self . leftBound) / 2
104+
105+ if leftBound <= middle { leftChild. update ( leftBound: leftBound, rightBound: rightBound, incremental: incremental) }
106+ if middle < rightBound { rightChild. update ( leftBound: leftBound, rightBound: rightBound, incremental: incremental) }
107+
108+ pushUp ( lson: leftChild, rson: rightChild)
109+ }
110+
111+ }
112+
113+ let array = [ 1 , 2 , 3 , 4 , 1 , 3 , 2 ]
114+
115+ let sumSegmentTree = LazySegmentTree ( array: array)
116+
117+ print ( sumSegmentTree. query ( leftBound: 0 , rightBound: 3 ) ) // 10 = 1 + 2 + 3 + 4
118+ sumSegmentTree. update ( index: 1 , incremental: 2 )
119+ print ( sumSegmentTree. query ( leftBound: 0 , rightBound: 3 ) ) // 12 = 1 + 4 + 3 + 4
120+ sumSegmentTree. update ( leftBound: 0 , rightBound: 2 , incremental: 2 )
121+ print ( sumSegmentTree. query ( leftBound: 0 , rightBound: 3 ) ) // 18 = 3 + 6 + 5 + 4
122+
123+ for index in 2 ... 5 {
124+ sumSegmentTree. update ( index: index, incremental: 3 )
125+ }
126+
127+
128+ sumSegmentTree. update ( leftBound: 0 , rightBound: 5 , incremental: 2 )
129+ print ( sumSegmentTree. query ( leftBound: 0 , rightBound: 2 ) )
0 commit comments