Ako sme spomínali, zmeníme celý spôsob implementácie intervalového stromu. Každý vrchol bude reprezentovaný objektom, ktorý si bude pamätať svoju hodnotu, lenivú hodnotu, začiatok a koniec intervalu, ktorého súčet jeho hodnota reprezentuje a referencie na jeho synov. Počiatočnú hodnotu vrcholu vieme v konštruktore jednoducho vypočítať, vďaka začiatku o koncu intervalu, ktorý daný vrchol reprezentuje. Začiatok a koniec intervalu nám tiež trochu zníži počet parametrov v intervalových operáciach, keďže už si nemusíme udržiavať tieto údaje ako parametre. Všimnime si, že intervalové operácie ostali nezmenené. Stačilo zmeniť metódu propagate, keďže ju voláme vždy predtým, ako pristúpime k synom aktuálneho vrcholu. Všetky intervalové operácie budeme volať z koreňu nášho stromu a to bude aj jediný vrchol, ktorý si v našom kóde budeme pamätať (mať uložený ako premennú).
class Node:
def __init__(self, i, j):
self.value = j - i
self.lazy = 1
self.i = i
self.j = j
self.left = None
self.right = None
def query(self, l, r):
if l >= self.j or r <= self.i:
return 0
if l <= self.i and r >= self.j:
return self.value
self.propagate()
return self.left.query(l, r) + self.right.query(l, r)
def update(self, l, r, x):
if l >= self.j or r <= self.i:
return
if l <= self.i and r >= self.j:
self.lazy *= x
self.value *= x
return
self.propagate()
self.left.update(l, r, x)
self.right.update(l, r, x)
self.value = self.left.value + self.right.value
def propagate(self):
mid = (self.i + self.j) // 2
if self.left == None:
self.left = Node(self.i, mid)
self.left.lazy *= self.lazy
self.left.value *= self.lazy
if self.right == None:
self.right = Node(mid, self.j)
self.right.lazy *= self.lazy
self.right.value *= self.lazy
self.lazy = 1