Tu samozrejme naša cesta nekončí. Keby sme len chceli usporiadanú množinu, mohli sme použiť štandardný knižničný set a mali by sme to bez práce. My ale vlastné stromy implementujeme práve preto, aby sme do nich mohli dorobiť funkcionalitu, ktorú nám štandardné knižnice neponúkajú.
Jednou často užitočnou fičúriou je prístup ku k-temu najmenšiemu prvku. Chceli by sme teda vedieť pre ľubovoľné $k$ vedieť v logaritmickom čase odpovedať na otázku: "keby sme všetky prvky z tohto stromu mali v usporiadanom poli, ktorý prvok by bol na indexe $k$?"
Takéto niečo vieme veľmi ľahko implementovať. Jediné, čo treba spraviť, je do každého vrcholu pridať premennú, v ktorej si pamätáme veľkosť podstromu, ktorého je koreňom. (Inými slovami, keby sme odstrihli hranu z tohto vrcholu do jeho otca, koľko vrcholov by odpadlo?)
Toto sa robí ľahko. Stačí vždy, keď nejaký vrchol meníme, prepočítať túto hodnotu z hodnôt jeho synov: veľkosť môjho podstromu je 1 (za mňa), plus veľkosť podstromu môjho ľavého syna (ak nejakého mám), plus veľkosť podstromu môjho pravého syna (detto).
Celá implementácia je teda len pár riadkov navyše. Začneme tým, že si vyrobíme dve pomocné funkcie:
int get_size(pitem t) { return t ? t->size : 0; }
void update(pitem t) { if (t) t->size = 1 + get_size(t->l) + get_size(t->r); }
No a teraz už len stačí pridať volanie update(t) na koniec každej z funkcií split, merge, insert a erase a máme strom, ktorý si správne udržiava veľkosti podstromov.
Akonáhle máme túto pomocnú informáciu, vieme v logaritmickom čase "indexovať" do nášho treapu. Predstavme si, že sme v koreni treapu a hľadáme prvok na indexe $k$. Pozrieme sa na size pre ľavého syna koreňa a dozvieme sa, že je $x$. Čo to znamená? Znamená to, že prvky na indexoch 0 až $x-1$ sú v ľavom podstrome koreňa, na indexe $x$ je koreň, a od indexu $x+1$ začínajú prvky v pravom podstrome koreňa. Podľa porovnania $k$ a $x$ vieme, kde želaný prvok hľadať (a ak ho
hľadáme v pravom podstrome, vieme, že v rámci neho chceme prvok s indexom $k-x-1$).
Implementácia:
T operator[](int x) { return kth_smallest(root,x); }
T kth_smallest(pitem p, int x) {
int l = get_size(p->l);
return l > x ? kth_smallest(p->l,x) : l == x ? p->key : kth_smallest(p->r,x-l-1);
}