avl.clj

(ns avl.avl
  (:refer-clojure :exclude [key find get val]))

(defprotocol Node
  (left   [node])
  (right  [node])
  (key    [node])
  (val    [node])
  (height [node])
  (add [tree key value])
  (-balance [node])
  (-rotate-l [node])
  (-rotate-r [node]))

(defn- tilt [n]
  (- (height (left n)) (height (right n))))

(defn find [node k]
  (when node
    (let [cmp (compare k (key node))]
      (cond
        (zero? cmp) (val node)
        (neg?  cmp) (recur (left  node) k)
        :else       (recur (right node) k)))))

(defrecord AVLNode [l r k v height]
  Node
  (-rotate-l [node]
    (let [l (->AVLNode l (left r) k v)]
      (->AVLNode l (right r) (key r) (val r))))
  (-rotate-r [node]
    (let [r (->AVLNode (right l) r k v)]
      (->AVLNode (left l) r (key l) (val l))))
  (-balance [node]
    (case (tilt node)
      +2 (-rotate-r
          (if (= 1 (tilt l))
            (->AVLNode (-rotate-l l) r k v)
            node))
      -2 (-rotate-l
          (if (= 1 (tilt r))
            (->AVLNode l (-rotate-r r) k v)
            node))
      node))
  (add [node k' v']
    (let [cmp (compare k' k)]
      (if (zero? cmp)
        (->AVLNode l r k v')
        (-balance
         (if (neg? cmp)
           (->AVLNode (add l k' v') r k v)
           (->AVLNode l (add r k' v') k v))))))
  (key    [node] k)
  (val    [node] v)
  (height [node] height)
  (left   [node] l)
  (right  [node] r))

(extend-type nil
  Node
  (add    [_ k v] (->AVLNode nil nil k v))
  (left   [_]     nil)
  (right  [_]     nil)
  (height [_]     0))

(defn ->AVLNode
  ([l r k v] (AVLNode. l r k v (inc (max (height l) (height r))))))