Implementácia

Existujú minimálne dva dobré spôsoby ako implemntovať intervalový strom:

• "V poli": celý strom si budeme pamätať v jednom poli. Výhodou tohto prístupu je kratší (a väčšinou rýchlejšie napísaný) kód. Nevýhodou je, že keď takýto intervalový strom chcete rozšíriť o ďalšiu funkcionalitu, kód začne byť nepríjemne zamotaný a ťažko čitateľný.

• Objektovo orientovaný: každý vrchol stromu bude jeden objekt a bude si pamätať referencie (smerníky) na svojich synov. Nevýhodou tohto spôsobu je trochu dlhší (a väčšinou pomalší) kód, výhodou je lepšia čitateľnosť kódu (ak sa kamarátite s objektovo orientovaným programovaním) a možnosť jednoduchého rozšírenia na perzistentný intervalový strom (čo je jedna zo silnejších verzií intervalového stromu).

My si ukážeme implemntáciu pomocou poľa (aj keď kvôli prehľadnosti budeme použivať python triedu na reprezentáciu celého intervalového stromu). Vrcholy nášho stromu si môžeme očíslovať po riadkoch zhora nadol číslami $1, 2, \dots, 2n-1$:

Toto číslovanie má jednu peknú vlastnosť: ak má vrchol číslo $x$, potom jeho synovia (ak nejakých má) majú čísla $2x$ a $2x+1$. Jeho otec má zasa číslo $\lfloor x/2 \rfloor$.

Vrcholy si teda môžeme pamätať v poli dĺžky $2n$, kde index $0$ nepoužijeme a inak bude na indexe $x$ vrchol číslo $x$. List číslo $i$ bude mať teda index $n+i$.

Vytvoriť intervalový strom, teda zapísať hodnoty do vrcholov vieme nasledovne. Najprv zapíšeme hodnoty daného poľa do posledných $n$ indexov poľa reprezentujúceho intervalový strom. Teraz, keď už vieme hodnoty listov, tak môžeme vypočítať hodnoty ostatných vrcholov. Ak pôjdeme odzadu (teda postupne po indexoch $n - 1, n - 2, \dots, 1$), takto máme stále dostatočnú informáciu na vypočítanie hodnoty aktuálne počítaného vrcholu. Takto vieme vytvoriť intervalový strom s lineárnou časovou zložitosťou od veľkosti poľa, podľa ktorého intervalový strom vytvárame. Je

Jednotlivé metódy teraz už vieme implementovať pomerne priamočiaro. Pri operácii súčtu intervalu je potrebné si pamätať okrem intervalu, ktorého súčet hľadáme aj interval, ktorého súčet aktuálny vrchol reprezentuje. To ale bez problémov dokážeme, keďže koreň reprezentuje interval celého stromu a pri volaní na nejakého zo synov vieme tento interval stále ľahko prepočítať.

import math

class SegmentTree:

    # vytvor intervalovy strom z daneho pola p
    def __init__(self, p):
        self.n = 2 ** math.ceil(math.log2(len(p)))
        self.values = [0] * (2 * self.n)
        for i in range(len(p)):
            self.values[i + self.n] = p[i]
        for i in range(self.n - 1, 0, -1):
            self.values[i] = self.values[i * 2] + self.values[i * 2 + 1]

    # sucet prvkov s indexami z [l, r) pod vrcholom v
    # i, j reprezentuju interval, ktoreho sucet mame zapisany vo vrchole v
    def query(self, v, i, j, l, r):
        # nic co je pod vrcholom v nelezi v intervale [l, r)
        if l >= j or r <= i:
            return 0
        # vsetko pod vrcholom v je v intervale [l, r)
        if l <= i and r >= j:
            return self.values[v]
        mid = (i + j) // 2
        # delegovanie na synov
        return self.query(v * 2, i, mid, l, r) + self.query(v * 2 + 1, mid, j, l, r)

    # zmen hodnotu p_i na x
    def update(self, i, x):
        v = i + self.n
        self.values[v] = x
        # postupujeme zdola nahor po predkoch v-teho listu
        while v > 1:
            v //= 2
            self.values[v] = self.values[v * 2] + self.values[v * 2 + 1]

# pouzitie
p = [1, 3, 1, 4]
segtree = SegmentTree(p)
print(segtree.query(1, 0, segtree.n, 2, 4))
segtree.update(3, 10)
print(segtree.query(1, 0, segtree.n, 2, 4))