-
Notifications
You must be signed in to change notification settings - Fork 0
/
SegmentTree.py
104 lines (81 loc) · 3.81 KB
/
SegmentTree.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import random
class Node:
def __init__(self, key, leftIdx, rightIdx):
self.key = key
self.leftChild = None
self.rightChild = None
# rango del nodo
self.leftIdx = leftIdx
self.rightIdx = rightIdx
class segmentTree:
def __init__(self, numbers):
self.root = Node(0, 0, len(numbers)-1)
self.numbers = numbers
def build(self, currentNode):
# cuando encontramos una hoja le asignamos el valor del index que apunta en la lista
if currentNode.leftIdx == currentNode.rightIdx:
currentNode.key = self.numbers[currentNode.leftIdx]
return
# Encontrar la mitad del rango, // para redondear a int
midIdx = (currentNode.leftIdx + currentNode.rightIdx)//2
# Creamos el nodo izq y der
currentNode.leftChild = Node(0, currentNode.leftIdx, midIdx)
currentNode.rightChild = Node(0, midIdx+1, currentNode.rightIdx)
# Creamos el subarabol izq y der
# llamamos la función build con el hijo izq y der
self.build(currentNode.leftChild)
self.build(currentNode.rightChild)
# Mezclamos los valores de los dos hijos en el nodo actuals
currentNode.key = currentNode.leftChild.key + currentNode.rightChild.key
def __build(self):
self.build(self.root)
def update(self, currentNode, addValue, idx):
# Si terminamos de recorrer la rama
if currentNode == None:
return
# Chechar si el rango del actual contiene al indicie a modifcar
if idx < currentNode.leftIdx or idx > currentNode.rightIdx:
return
# Modificar el valor del nodo
# Hacer llamada recurisva para modificiar el valor de sus hijos
currentNode.key += addValue
self.update(currentNode.leftChild, addValue, idx)
self.update(currentNode.rightChild, addValue, idx)
def update2(self, currentNode, newValue, idx):
# Si terminamos de recorrer la rama
if currentNode == None:
return
# Chechar si el rango del actual contiene al indicie a modifcar
if idx < currentNode.leftIdx or idx > currentNode.rightIdx:
return
if currentNode.leftIdx == currentNode.rightIdx:
currentNode.key = newValue
return
# Modificar el valor del nodo
# Hacer llamada recurisva para modificiar el valor de sus hijos
self.update2(currentNode.leftChild, newValue, idx)
self.update2(currentNode.rightChild, newValue, idx)
currentNode.key = currentNode.leftChild.key + currentNode.rightChild.key
def getSum(self, currentNode, l, r):
if currentNode is None:
return 0
# si el indice del nodo este fuera del rango no sumamos nada
if r < currentNode.leftIdx or l > currentNode.rightIdx:
return 0
# si el indice del nodo está dentro del rango retornamos el valor del nodo
elif l <= currentNode.leftIdx and r >= currentNode.rightIdx:
return currentNode.key
# si el nodo está parcialmente dentro del rango, volvermos a llamar la función tanto con el hijo izquierdo como el derecho
else:
return self.getSum(currentNode.leftChild, l, r)+self.getSum(currentNode.rightChild, l, r)
numbers = [5, 4, 3, 2, 10, 0, 3, 4]
print("Arreglo orginal: ", numbers)
st = segmentTree(numbers)
st.build(st.root)
print("Obtenemos suma del 2-5: ", st.getSum(st.root, 2, 5))
st.update(st.root, 6, 4)
print("\nSumanos 6 al nodo 4\nVolvemos a obtener la suma del 2-5: ",
st.getSum(st.root, 2, 5))
st.update2(st.root, 0, 3)
print("\nCambiamos al nodo 3 dandole el valor de 0\nVolvemos a obtener la suma del 2-5: ",
st.getSum(st.root, 2, 5), "\n")