Treaps
June 26, 2009
We represent a treap as a tree of nodes, each a five-slot vector with places for the key, priority, value, and the two child nodes; access functions are given below:
(define (treap k p v l r) (vector k p v l r))
(define (key t) (vector-ref t 0))
(define (prio t) (vector-ref t 1))
(define (val t) (vector-ref t 2))
(define (lkid t) (vector-ref t 3))
(define (rkid t) (vector-ref t 4))
(define nil (vector 'nil -1 'nil 'nil 'nil))
(define (nil! k) (vector-set! nil 0 k))
(define (nil? t) (eq? t nil))
(define (leaf? t) (and (nil? (lkid t)) (nil? (rkid t))))
(define (leaf-or-nil? t) (eq? (lkid t) (rkid t)))
Rotations are done by building new nodes using the appropriate pieces of the old nodes:
(define (rot-left t)
(let ((l (treap (key t) (prio t) (val t) (lkid t) (lkid (rkid t)))))
(treap (key (rkid t)) (prio (rkid t)) (val (rkid t)) l (rkid (rkid t)))))
(define (rot-right t)
(let ((r (treap (key t) (prio t) (val t) (rkid (lkid t)) (rkid t))))
(treap (key (lkid t)) (prio (lkid t)) (val (lkid t)) (lkid (lkid t)) r)))
Lookup is exactly the same as a standard binary search tree; the priority field is unused. Starting at the root, lookup compares the key being sought to the current key, branching either left or right depending on the outcome of the comparison. Lookup terminates with failure if it reaches nil, and terminates with success if the compare returns equal.
(define (lookup lt? t k)
(cond ((nil? t) #f)
((lt? k (key t)) (lookup lt? (lkid t) k))
((lt? (key t) k) (lookup lt? (rkid t) k))
(else (cons k (val t)))))
Insert works in two phases; a winding phase descends the tree to find the insertion point, in a manner similar to lookup, and an unwinding phase performs rotations to restore the heap-order property:
(define (insert lt? t k v)
(cond ((nil? t) (treap k (rand) v nil nil))
((lt? k (key t))
(let ((t (treap (key t) (prio t) (val t) (insert lt? (lkid t) k v) (rkid t))))
(if (< (prio t) (prio (lkid t))) (rot-right t) t)))
((lt? (key t) k)
(let ((t (treap (key t) (prio t) (val t) (lkid t) (insert lt? (rkid t) k v))))
(if (< (prio t) (prio (rkid t))) (rot-left t) t)))
(else (treap k (prio t) v (lkid t) (rkid t)))))
Delete begins by setting the desired key in nil as a sentinel, then searches, rebuilding nodes as it proceeds, until the key being deleted is found (the search will always succeed because of the sentinel). Then, delete calls an auxiliary procedure, deroot, which rotates the current node down until it becomes a leaf, where it is clipped off. Deroot selects rotations based on priority:
(define (deroot t)
(cond ((leaf-or-nil? t) nil)
((< (prio (lkid t)) (prio (rkid t)))
(let ((t (rot-left t)))
(treap (key t) (prio t) (val t) (deroot (lkid t)) (rkid t))))
(else (let ((t (rot-right t)))
(treap (key t) (prio t) (val t) (lkid t) (deroot (rkid t)))))))
(define (delete lt? t k)
(nil! k)
(let delete ((t t))
(cond ((lt? k (key t))
(treap (key t) (prio t) (val t) (delete (lkid t)) (rkid t)))
((lt? (key t) k)
(treap (key t) (prio t) (val t) (lkid t) (delete (rkid t))))
(else (deroot t)))))
Update is a simple variant on insert:
(define (update lt? t f k v)
(cond ((nil? t) (treap k (rand) v nil nil))
((lt? k (key t))
(let ((t (treap (key t) (prio t) (val t) (update lt? (lkid t) k (rand) v) (rkid t))))
(if (< (prio t) (prio (lkid t))) (rot-right t) t)))
((lt? (key t) k)
(let ((t (treap (key t) (prio t) (val t) (lkid t) (update lt? (rkid t) k (rand) v))))
(if (< (prio t) (prio (rkid t))) (rot-left t) t)))
(else (treap k (prio t) (f k (val t)) (lkid t) (rkid t)))))
Enlist traverses the tree in order, calling itself recursively at each node:
(define (enlist t)
(if (nil? t) '()
(append (enlist (lkid t)) (list (cons (key t) (val t))) (enlist (rkid t)))))
An example of the use of treaps is given by the earlier word-frequency exercise, which counted the most frequently-used words in the Bible. Here is a restatement of that exercise using treaps:
(define (word-freq n file-name)
(define (freq-gt? a b) (> (cdr a) (cdr b)))
(with-input-from-file file-name
(lambda ()
(let loop ((word (read-word)) (freqs nil))
(if (eof-object? word)
(take n (sort freq-gt? (enlist freqs)))
(loop (read-word) (update string<? freqs (lambda (k v) (+ v 1)) word 1)))))))
The result is the same as the prior exercise:
> (word-freq 25 "bible.txt")
(("the" . 62588) ("and" . 30875) ("of" . 30183)
("to" . 23023) ("you" . 14887) ("in" . 13357) ("he" . 10495)
("a" . 10150) ("i" . 9078) ("for" . 8983) ("his" . 8424)
("lord" . 8129) ("your" . 7398) ("with" . 7259)
("that" . 7187) ("is" . 7143) ("they" . 7005) ("not" . 6484)
("him" . 6140) ("will" . 6093) ("them" . 5831) ("be" . 5668)
("who" . 5611) ("from" . 5476) ("it" . 5395))
Rand comes from the Standard Prelude. Treaps were invented by Cecilia Aragon and Raimund Seidel, and described in two papers. You can see the code from this exercise at http://programmingpraxis.codepad.org/iFm0jW0A.
Programming Praxis 
[…] Praxis – Treaps By Remco Niemeijer Today’s Programming Praxis problem is about Treaps – binary trees that are more or less balanced thanks […]
My Haskell solution (see http://bonsaicode.wordpress.com/2009/06/26/programming-praxis-treaps/ for a version with comments):
import Control.Monad import Data.Char import qualified Data.List.Key as K import System.Random data Treap k a = Nil | Node Int k a (Treap k a) (Treap k a) priority :: Treap k a -> Int priority Nil = -1 priority (Node p _ _ _ _) = p rotLeft :: Treap k a -> Treap k a rotLeft (Node p k a l (Node rp rk ra rl rr)) = Node rp rk ra (Node p k a l rl) rr rotLeft t = t rotRight :: Treap k a -> Treap k a rotRight (Node p k a (Node lp lk la ll lr) r) = Node lp lk la ll (Node p k a lr r) rotRight t = t rot :: Treap k a -> Treap k a rot Nil = Nil rot t@(Node p _ _ l r) | p < priority l = rotRight t | p < priority r = rotLeft t | otherwise = t find :: Ord k => k -> Treap k a -> Maybe a find _ Nil = Nothing find k' (Node _ k a l r) | k' < k = find k' l | k' > k = find k' r | otherwise = Just a update :: Ord k => (a -> a -> a) -> k -> a -> Treap k a -> IO (Treap k a) update _ k' a' Nil = fmap (\r -> Node r k' a' Nil Nil) $ randomRIO (0, maxBound) update f k' a' (Node p k a l r) | k' < k = fmap (\n -> rot $ Node p k a n r) (update f k' a' l) | k' > k = fmap (rot . Node p k a l) (update f k' a' r) | otherwise = return $ Node p k' (f a' a) l r insert :: Ord k => k -> a -> Treap k a -> IO (Treap k a) insert = update const deroot :: Treap k a -> Treap k a deroot Nil = Nil deroot t@(Node _ _ _ l r) | priority l < priority r = d deroot id $ rotLeft t | otherwise = d id deroot $ rotRight t where d fl fr = (\(Node p k a l' r') -> Node p k a (fl l') (fr r')) delete :: Ord k => k -> Treap k a -> Treap k a delete _ Nil = Nil delete k' t@(Node p k a l r) | k' < k = Node p k a (delete k' l) r | k' > k = Node p k a l (delete k' r) | otherwise = deroot t toList :: Treap k a -> [(k, a)] toList Nil = [] toList (Node _ k a l r) = toList l ++ [(k, a)] ++ toList r main :: IO () main = mapM_ print =<< wordFreqs 25 =<< readFile "bible.txt" wordFreqs :: Int -> String -> IO [(String, Int)] wordFreqs n = fmap (take n . reverse . K.sort snd . toList) . foldM (\a w -> update (+) w 1 a) Nil . map (filter isAlpha) . wordsAttempt in clojure , using pseudocode from http://sims.berkeley.edu/~aragon/pubs/rst96.pdf
pastebin link http://clojure.pastebin.com/f20e92d3a
(defstruct treaps :key :prio :value :left :right) (defn less? [a b] (neg? (compare a b))) (defn greater? [a b] (pos? (compare a b))) (defn eql? [a b] (zero? (compare a b))) (defn insert [k p v treap] (cond (nil? treap) (struct-map treaps :key k :prio p :value v) (less? k (:key treap)) (let [tr (assoc treap :left (insert k p v (:left treap)))] (if (> (:prio (:left tr)) (:prio tr)) (rotate-right tr) tr)) (greater? k (:key treap)) (let [tr (assoc treap :right (insert k p v (:right treap)))] (if (> (:prio (:right tr)) (:prio tr)) (rotate-left tr) tr)) :else treap)) (defn delete [k treap] (cond (nil? treap) nil (less? k (:key treap)) (assoc treap :left (delete k (:left treap))) (greater? k (:key treap)) (assoc treap :right (delete k (:right treap))) :else (root-delete treap))) (defn root-delete [treap] (cond (and (nil? (:left treap)) (nil? (:right treap))) nil (nil? (:left treap)) (:right treap) (nil? (:right treap)) (:left treap) :else (if (> (:prio (:left treap)) (:prio (:right treap))) (let [tr (rotate-right treap)] (assoc tr :right (root-delete (:right tr)))) (let [tr (rotate-left treap)] (assoc tr :left (root-delete (:left tr))))))) (defn look-up [k treap] (cond (nil? treap) nil (less? k (:key treap)) (look-up k (:left treap)) (greater? k (:key treap)) (look-up k (:right treap)) :else (:value treap))) (defn enlist [treap] (if (nil? treap) '() (concat (enlist (:left treap)) (list (str (:key treap)) (:prio treap)) (enlist (:right treap))))) (defn update [k v treap] (cond (nil? treap) nil (less? k (:key treap)) (assoc treap :left (update k v (:left treap))) (greater? k (:key treap)) (assoc treap :right (update k v (:right treap))) :else (assoc treap :value v))) (defn rotate-right [a-tree] (assoc (:left a-tree) :right (assoc a-tree :left (:right (:left a-tree))))) (defn rotate-left [a-tree] (assoc (:right a-tree) :left (assoc a-tree :right (:left (:right a-tree))))) (def lst '((80 \v) (60 \g) (63 \z) (37 \a) (57 \s) (47 \x) (31 \d) (53 \k) (39 \u) (22 \w) (36 \y) (15 \j) (48 \p) (21 \t) (17 \m) (34 \q) )) (def tree (reduce (fn [a b] (insert (second b) (first b) (first b) a)) nil lst)) (def tree2 (reduce (fn [a b] (insert (second b) (rand) (first b) a)) nil lst)) (= (delete \l (insert \l 69 69 tree)) tree) (= (delete \l (insert \l (rand) "L" tree2)) tree2)Finally finished. It took me a while to get everything working since I went all out and implemented the full IDictionary interface. It was fun to get all of the code to work properly, but there’s clearly some optimization work needed to make the thing useful. My implementation’s performance compares unfavorably to the .Net framework’s built in System.Collections.Generic.Dictionary in all of my tests.
http://inscrutable.pastebin.com/f3fabfc9a
My implementation in C
#include <stdlib.h> #include <time.h> typedef struct treapNode { int key; int priority; struct treapNode* left; struct treapNode* right; }treap; static treap* sentinel; static treap* create_node(int key); static treap* right_rotation(treap *p); static treap* left_rotation(treap *p); treap* treap_init() { sentinel = (treap*)malloc(sizeof(*sentinel)); if(sentinel == NULL) return NULL; sentinel->key = 0; sentinel->priority = INT_MIN; sentinel->left = sentinel->right = NULL; srand((unsigned int)time(NULL)); return sentinel; } treap* treap_search(treap *p, int key) { if(p == sentinel) return NULL; else if(key == p->key) return p; else if(key < p->key) return treap_search(p->left, key); else return treap_search(p->right, key); } treap* treap_insert(treap *p, int key) { if(p == sentinel) { p = create_node(key); if(p == NULL) return NULL; } else if(key < p->key) { p->left = treap_insert(p->left, key); if(p->priority < p->left->priority) p = right_rotation(p); } else { p->right = treap_insert(p->right, key); if(p->priority < p->right->priority) p = left_rotation(p); } return p; } treap* treap_remove(treap *p, int key) { if(key < p->key) p->left = treap_remove(p->left, key); else if(key > p->key) p->right = treap_remove(p->right, key); else { if(p->left == sentinel && p->right == sentinel) { free(p); return sentinel; } if(p->left->priority > p->right->priority) { p = right_rotation(p); p->right = treap_remove(p->right, key); } else { p = left_rotation(p); p->left = treap_remove(p->left, key); } } return p; } static treap* create_node(int key) { treap *p; p = (treap*)malloc(sizeof(*p)); if(p == NULL) return NULL; p->priority = rand(); p->key = key; p->left = p->right = sentinel; return p; } static treap* right_rotation(treap *p) { treap *s; s = p->left; p->left = s->right; s->right = p; return s; } static treap* left_rotation(treap *p) { treap *s; s = p->right; p->right = s->left; s->left = p; return s; }