iframe-proxy | Sunbelt Computer Software

History

241 lines (220 loc) · 6.54 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

"""

Segment_tree creates a segment tree with a given array and function,

allowing queries to be done later in log(N) time

function takes 2 values and returns a same type value

"""

from collections.abc import Sequence

from queue import Queue

class SegmentTreeNode:

def __init__(self, start, end, val, left=None, right=None):

self.start = start

self.end = end

self.val = val

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

self.left = left

self.right = right

def __str__(self):

return f"val: {self.val}, start: {self.start}, end: {self.end}"

class SegmentTree:

"""

>>> import operator

>>> num_arr = SegmentTree([2, 1, 5, 3, 4], operator.add)

>>> for node in num_arr.traverse():

... print(node)

...

val: 15, start: 0, end: 4

val: 8, start: 0, end: 2

val: 7, start: 3, end: 4

val: 3, start: 0, end: 1

val: 5, start: 2, end: 2

val: 3, start: 3, end: 3

val: 4, start: 4, end: 4

val: 2, start: 0, end: 0

val: 1, start: 1, end: 1

>>>

>>> num_arr.update(1, 5)

>>> for node in num_arr.traverse():

... print(node)

...

val: 19, start: 0, end: 4

val: 12, start: 0, end: 2

val: 7, start: 3, end: 4

val: 7, start: 0, end: 1

val: 5, start: 2, end: 2

val: 3, start: 3, end: 3

val: 4, start: 4, end: 4

val: 2, start: 0, end: 0

val: 5, start: 1, end: 1

>>>

>>> num_arr.query_range(3, 4)

>>> num_arr.query_range(2, 2)

>>> num_arr.query_range(1, 3)

>>>

>>> max_arr = SegmentTree([2, 1, 5, 3, 4], max)

>>> for node in max_arr.traverse():

... print(node)

...

val: 5, start: 0, end: 4

val: 5, start: 0, end: 2

val: 4, start: 3, end: 4

val: 2, start: 0, end: 1

val: 5, start: 2, end: 2

val: 3, start: 3, end: 3

val: 4, start: 4, end: 4

val: 2, start: 0, end: 0

val: 1, start: 1, end: 1

>>>

>>> max_arr.update(1, 5)

>>> for node in max_arr.traverse():

... print(node)

...

val: 5, start: 0, end: 4

val: 5, start: 0, end: 2

val: 4, start: 3, end: 4

val: 5, start: 0, end: 1

val: 5, start: 2, end: 2

val: 3, start: 3, end: 3

val: 4, start: 4, end: 4

val: 2, start: 0, end: 0

val: 5, start: 1, end: 1

>>>

>>> max_arr.query_range(3, 4)

>>> max_arr.query_range(2, 2)

>>> max_arr.query_range(1, 3)

>>>

>>> min_arr = SegmentTree([2, 1, 5, 3, 4], min)

>>> for node in min_arr.traverse():

... print(node)

...

val: 1, start: 0, end: 4

val: 1, start: 0, end: 2

val: 3, start: 3, end: 4

val: 1, start: 0, end: 1

val: 5, start: 2, end: 2

val: 3, start: 3, end: 3

val: 4, start: 4, end: 4

val: 2, start: 0, end: 0

val: 1, start: 1, end: 1

>>>

>>> min_arr.update(1, 5)

>>> for node in min_arr.traverse():

... print(node)

...

val: 2, start: 0, end: 4

val: 2, start: 0, end: 2

val: 3, start: 3, end: 4

val: 2, start: 0, end: 1

val: 5, start: 2, end: 2

val: 3, start: 3, end: 3

val: 4, start: 4, end: 4

val: 2, start: 0, end: 0

val: 5, start: 1, end: 1

>>>

>>> min_arr.query_range(3, 4)

>>> min_arr.query_range(2, 2)

>>> min_arr.query_range(1, 3)

>>>

"""

def __init__(self, collection: Sequence, function):

self.collection = collection

self.fn = function

if self.collection:

self.root = self._build_tree(0, len(collection) - 1)

def update(self, i, val):

"""

Update an element in log(N) time

:param i: position to be update

:param val: new value

>>> import operator

>>> num_arr = SegmentTree([2, 1, 5, 3, 4], operator.add)

>>> num_arr.update(1, 5)

>>> num_arr.query_range(1, 3)

"""

self._update_tree(self.root, i, val)

def query_range(self, i, j):

"""

Get range query value in log(N) time

:param i: left element index

:param j: right element index

:return: element combined in the range [i, j]

>>> import operator

>>> num_arr = SegmentTree([2, 1, 5, 3, 4], operator.add)

>>> num_arr.update(1, 5)

>>> num_arr.query_range(3, 4)

>>> num_arr.query_range(2, 2)

>>> num_arr.query_range(1, 3)

>>>

"""

return self._query_range(self.root, i, j)

def _build_tree(self, start, end):

if start == end:

return SegmentTreeNode(start, end, self.collection[start])

mid = (start + end) // 2

left = self._build_tree(start, mid)

right = self._build_tree(mid + 1, end)

return SegmentTreeNode(start, end, self.fn(left.val, right.val), left, right)

def _update_tree(self, node, i, val):

if node.start == i and node.end == i:

node.val = val

return

if i <= node.mid:

self._update_tree(node.left, i, val)

else:

self._update_tree(node.right, i, val)

node.val = self.fn(node.left.val, node.right.val)

def _query_range(self, node, i, j):

if node.start == i and node.end == j:

return node.val

if i <= node.mid:

if j <= node.mid:

# range in left child tree

return self._query_range(node.left, i, j)

else:

# range in left child tree and right child tree

return self.fn(

self._query_range(node.left, i, node.mid),

self._query_range(node.right, node.mid + 1, j),

)

else:

# range in right child tree

return self._query_range(node.right, i, j)

def traverse(self):

if self.root is not None:

queue = Queue()

queue.put(self.root)

while not queue.empty():

node = queue.get()

yield node

if node.left is not None:

queue.put(node.left)

if node.right is not None:

queue.put(node.right)

if __name__ == "__main__":

import operator

for fn in [operator.add, max, min]:

print("*" * 50)

arr = SegmentTree([2, 1, 5, 3, 4], fn)

for node in arr.traverse():

print(node)

print()

arr.update(1, 5)

for node in arr.traverse():

print(node)

print()

print(arr.query_range(3, 4)) # 7

print(arr.query_range(2, 2)) # 5

print(arr.query_range(1, 3)) # 13

print()

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Sunbelt Computer Software

PL/B Language Development and Support

FilesExpand file tree

segment_tree_other.py

Latest commit

History

segment_tree_other.py

File metadata and controls