iframe-proxy | Sunbelt Computer Software

History

116 lines (100 loc) · 3.2 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

import math

class SegmentTree:

def __init__(self, a):

self.A = a

self.N = len(self.A)

self.st = [0] * (

4 * self.N

) # approximate the overall size of segment tree with array N

if self.N:

self.build(1, 0, self.N - 1)

def left(self, idx):

"""

Returns the left child index for a given index in a binary tree.

>>> s = SegmentTree([1, 2, 3])

>>> s.left(1)

>>> s.left(2)

"""

return idx * 2

def right(self, idx):

"""

Returns the right child index for a given index in a binary tree.

>>> s = SegmentTree([1, 2, 3])

>>> s.right(1)

>>> s.right(2)

"""

return idx * 2 + 1

def build(self, idx, left, right):

if left == right:

self.st[idx] = self.A[left]

else:

mid = (left + right) // 2

self.build(self.left(idx), left, mid)

self.build(self.right(idx), mid + 1, right)

self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)])

def update(self, a, b, val):

"""

Update the values in the segment tree in the range [a,b] with the given value.

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

>>> s.update(2, 4, 10)

True

>>> s.query(1, 5)

"""

return self.update_recursive(1, 0, self.N - 1, a - 1, b - 1, val)

def update_recursive(self, idx, left, right, a, b, val):

"""

update(1, 1, N, a, b, v) for update val v to [a,b]

"""

if right < a or left > b:

return True

if left == right:

self.st[idx] = val

return True

mid = (left + right) // 2

self.update_recursive(self.left(idx), left, mid, a, b, val)

self.update_recursive(self.right(idx), mid + 1, right, a, b, val)

self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)])

return True

def query(self, a, b):

"""

Query the maximum value in the range [a,b].

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

>>> s.query(1, 3)

>>> s.query(1, 5)

"""

return self.query_recursive(1, 0, self.N - 1, a - 1, b - 1)

def query_recursive(self, idx, left, right, a, b):

"""

query(1, 1, N, a, b) for query max of [a,b]

"""

if right < a or left > b:

return -math.inf

if left >= a and right <= b:

return self.st[idx]

mid = (left + right) // 2

q1 = self.query_recursive(self.left(idx), left, mid, a, b)

q2 = self.query_recursive(self.right(idx), mid + 1, right, a, b)

return max(q1, q2)

def show_data(self):

show_list = []

for i in range(1, self.N + 1):

show_list += [self.query(i, i)]

print(show_list)

if __name__ == "__main__":

A = [1, 2, -4, 7, 3, -5, 6, 11, -20, 9, 14, 15, 5, 2, -8]

N = 15

segt = SegmentTree(A)

print(segt.query(4, 6))

print(segt.query(7, 11))

print(segt.query(7, 12))

segt.update(1, 3, 111)

print(segt.query(1, 15))

segt.update(7, 8, 235)

segt.show_data()

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.py

Latest commit

History

segment_tree.py

File metadata and controls