lisp/emacs-lisp/avl-tree.el

   1 ;;; avl-tree.el --- balanced binary trees, AVL-trees
   2
   3 ;; Copyright (C) 1995, 2007, 2008  Free Software Foundation, Inc.
   4
   5 ;; Author: Per Cederqvist <ceder@lysator.liu.se>
   6 ;;      Inge Wallin <inge@lysator.liu.se>
   7 ;;      Thomas Bellman <bellman@lysator.liu.se>
   8 ;; Maintainer: FSF
   9 ;; Created: 10 May 1991
  10 ;; Keywords: extensions, data structures
  11
  12 ;; This file is part of GNU Emacs.
  13
  14 ;; GNU Emacs is free software: you can redistribute it and/or modify
  15 ;; it under the terms of the GNU General Public License as published by
  16 ;; the Free Software Foundation, either version 3 of the License, or
  17 ;; (at your option) any later version.
  18
  19 ;; GNU Emacs is distributed in the hope that it will be useful,
  20 ;; but WITHOUT ANY WARRANTY; without even the implied warranty of
  21 ;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  22 ;; GNU General Public License for more details.
  23
  24 ;; You should have received a copy of the GNU General Public License
  25 ;; along with GNU Emacs.  If not, see <http://www.gnu.org/licenses/>.
  26
  27 ;;; Commentary:
  28
  29 ;; An AVL tree is a nearly-perfect balanced binary tree.  A tree consists of
  30 ;; two elements, the root node and the compare function.  The actual tree
  31 ;; has a dummy node as its root with the real root in the left pointer.
  32 ;;
  33 ;; Each node of the tree consists of one data element, one left
  34 ;; sub-tree and one right sub-tree.  Each node also has a balance
  35 ;; count, which is the difference in depth of the left and right
  36 ;; sub-trees.
  37 ;;
  38 ;; The functions with names of the form "avl-tree--" are intended for
  39 ;; internal use only.
  40
  41 ;;; Code:
  42
  43 (eval-when-compile (require 'cl))
  44
  45 ;; ================================================================
  46 ;;; Functions and macros handling an AVL tree node.
  47
  48 (defstruct (avl-tree--node
  49             ;; We force a representation without tag so it matches the
  50             ;; pre-defstruct representation.  Also we use the underlying
  51             ;; representation in the implementation of avl-tree--node-branch.
  52             (:type vector)
  53             (:constructor nil)
  54             (:constructor avl-tree--node-create (left right data balance))
  55             (:copier nil))
  56   left right data balance)
  57
  58 (defalias 'avl-tree--node-branch 'aref
  59   ;; This implementation is efficient but breaks the defstruct abstraction.
  60   ;; An alternative could be
  61   ;; (funcall (aref [avl-tree-left avl-tree-right avl-tree-data] branch) node)
  62   "Get value of a branch of a node.
  63
  64 NODE is the node, and BRANCH is the branch.
  65 0 for left pointer, 1 for right pointer and 2 for the data.\"
  66 \(fn node branch)")
  67 ;; The funcall/aref trick doesn't work for the setf method, unless we try
  68 ;; and access the underlying setter function, but this wouldn't be
  69 ;; portable either.
  70 (defsetf avl-tree--node-branch aset)
  71
  72 \f
  73 ;; ================================================================
  74 ;;; Internal functions for use in the AVL tree package
  75
  76 (defstruct (avl-tree-
  77             ;; A tagged list is the pre-defstruct representation.
  78             ;; (:type list)
  79             :named
  80             (:constructor nil)
  81             (:constructor avl-tree-create (cmpfun))
  82             (:predicate avl-tree-p)
  83             (:copier nil))
  84   (dummyroot (avl-tree--node-create nil nil nil 0))
  85   cmpfun)
  86
  87 (defmacro avl-tree--root (tree)
  88   ;; Return the root node for an avl-tree.  INTERNAL USE ONLY.
  89   `(avl-tree--node-left (avl-tree--dummyroot tree)))
  90 (defsetf avl-tree--root (tree) (node)
  91   `(setf (avl-tree--node-left (avl-tree--dummyroot ,tree)) ,node))
  92
  93 ;; ----------------------------------------------------------------
  94 ;;                          Deleting data
  95
  96 (defun avl-tree--del-balance1 (node branch)
  97   ;; Rebalance a tree and return t if the height of the tree has shrunk.
  98   (let ((br (avl-tree--node-branch node branch))
  99         p1 b1 p2 b2 result)
 100     (cond
 101      ((< (avl-tree--node-balance br) 0)
 102       (setf (avl-tree--node-balance br) 0)
 103       t)
 104
 105      ((= (avl-tree--node-balance br) 0)
 106       (setf (avl-tree--node-balance br) +1)
 107       nil)
 108
 109      (t
 110       ;; Rebalance.
 111       (setq p1 (avl-tree--node-right br)
 112             b1 (avl-tree--node-balance p1))
 113       (if (>= b1 0)
 114           ;; Single RR rotation.
 115           (progn
 116             (setf (avl-tree--node-right br) (avl-tree--node-left p1))
 117             (setf (avl-tree--node-left p1) br)
 118             (if (= 0 b1)
 119                 (progn
 120                   (setf (avl-tree--node-balance br) +1)
 121                   (setf (avl-tree--node-balance p1) -1)
 122                   (setq result nil))
 123               (setf (avl-tree--node-balance br) 0)
 124               (setf (avl-tree--node-balance p1) 0)
 125               (setq result t))
 126             (setf (avl-tree--node-branch node branch) p1)
 127             result)
 128
 129         ;; Double RL rotation.
 130         (setq p2 (avl-tree--node-left p1)
 131               b2 (avl-tree--node-balance p2))
 132         (setf (avl-tree--node-left p1) (avl-tree--node-right p2))
 133         (setf (avl-tree--node-right p2) p1)
 134         (setf (avl-tree--node-right br) (avl-tree--node-left p2))
 135         (setf (avl-tree--node-left p2) br)
 136         (setf (avl-tree--node-balance br) (if (> b2 0) -1 0))
 137         (setf (avl-tree--node-balance p1) (if (< b2 0) +1 0))
 138         (setf (avl-tree--node-branch node branch) p2)
 139         (setf (avl-tree--node-balance p2) 0)
 140         t)))))
 141
 142 (defun avl-tree--del-balance2 (node branch)
 143   (let ((br (avl-tree--node-branch node branch))
 144         p1 b1 p2 b2 result)
 145     (cond
 146      ((> (avl-tree--node-balance br) 0)
 147       (setf (avl-tree--node-balance br) 0)
 148       t)
 149
 150      ((= (avl-tree--node-balance br) 0)
 151       (setf (avl-tree--node-balance br) -1)
 152       nil)
 153
 154      (t
 155       ;; Rebalance.
 156       (setq p1 (avl-tree--node-left br)
 157             b1 (avl-tree--node-balance p1))
 158       (if (<= b1 0)
 159           ;; Single LL rotation.
 160           (progn
 161             (setf (avl-tree--node-left br) (avl-tree--node-right p1))
 162             (setf (avl-tree--node-right p1) br)
 163             (if (= 0 b1)
 164                 (progn
 165                   (setf (avl-tree--node-balance br) -1)
 166                   (setf (avl-tree--node-balance p1) +1)
 167                   (setq result nil))
 168               (setf (avl-tree--node-balance br) 0)
 169               (setf (avl-tree--node-balance p1) 0)
 170               (setq result t))
 171             (setf (avl-tree--node-branch node branch) p1)
 172             result)
 173
 174         ;; Double LR rotation.
 175         (setq p2 (avl-tree--node-right p1)
 176               b2 (avl-tree--node-balance p2))
 177         (setf (avl-tree--node-right p1) (avl-tree--node-left p2))
 178         (setf (avl-tree--node-left p2) p1)
 179         (setf (avl-tree--node-left br) (avl-tree--node-right p2))
 180         (setf (avl-tree--node-right p2) br)
 181         (setf (avl-tree--node-balance br) (if (< b2 0) +1 0))
 182         (setf (avl-tree--node-balance p1) (if (> b2 0) -1 0))
 183         (setf (avl-tree--node-branch node branch) p2)
 184         (setf (avl-tree--node-balance p2) 0)
 185         t)))))
 186
 187 (defun avl-tree--do-del-internal (node branch q)
 188   (let ((br (avl-tree--node-branch node branch)))
 189     (if (avl-tree--node-right br)
 190         (if (avl-tree--do-del-internal br +1 q)
 191             (avl-tree--del-balance2 node branch))
 192       (setf (avl-tree--node-data q) (avl-tree--node-data br))
 193       (setf (avl-tree--node-branch node branch)
 194             (avl-tree--node-left br))
 195       t)))
 196
 197 (defun avl-tree--do-delete (cmpfun root branch data)
 198   ;; Return t if the height of the tree has shrunk.
 199   (let ((br (avl-tree--node-branch root branch)))
 200     (cond
 201      ((null br)
 202       nil)
 203
 204      ((funcall cmpfun data (avl-tree--node-data br))
 205       (if (avl-tree--do-delete cmpfun br 0 data)
 206           (avl-tree--del-balance1 root branch)))
 207
 208      ((funcall cmpfun (avl-tree--node-data br) data)
 209       (if (avl-tree--do-delete cmpfun br 1 data)
 210           (avl-tree--del-balance2 root branch)))
 211
 212      (t
 213       ;; Found it.  Let's delete it.
 214       (cond
 215        ((null (avl-tree--node-right br))
 216         (setf (avl-tree--node-branch root branch) (avl-tree--node-left br))
 217         t)
 218
 219        ((null (avl-tree--node-left br))
 220         (setf (avl-tree--node-branch root branch) (avl-tree--node-right br))
 221         t)
 222
 223        (t
 224         (if (avl-tree--do-del-internal br 0 br)
 225             (avl-tree--del-balance1 root branch))))))))
 226
 227 ;; ----------------------------------------------------------------
 228 ;;                           Entering data
 229
 230 (defun avl-tree--enter-balance1 (node branch)
 231   ;; Rebalance a tree and return t if the height of the tree has grown.
 232   (let ((br (avl-tree--node-branch node branch))
 233         p1 p2 b2 result)
 234     (cond
 235      ((< (avl-tree--node-balance br) 0)
 236       (setf (avl-tree--node-balance br) 0)
 237       nil)
 238
 239      ((= (avl-tree--node-balance br) 0)
 240       (setf (avl-tree--node-balance br) +1)
 241       t)
 242
 243      (t
 244       ;; Tree has grown => Rebalance.
 245       (setq p1 (avl-tree--node-right br))
 246       (if (> (avl-tree--node-balance p1) 0)
 247           ;; Single RR rotation.
 248           (progn
 249             (setf (avl-tree--node-right br) (avl-tree--node-left p1))
 250             (setf (avl-tree--node-left p1) br)
 251             (setf (avl-tree--node-balance br) 0)
 252             (setf (avl-tree--node-branch node branch) p1))
 253
 254         ;; Double RL rotation.
 255         (setq p2 (avl-tree--node-left p1)
 256               b2 (avl-tree--node-balance p2))
 257         (setf (avl-tree--node-left p1) (avl-tree--node-right p2))
 258         (setf (avl-tree--node-right p2) p1)
 259         (setf (avl-tree--node-right br) (avl-tree--node-left p2))
 260         (setf (avl-tree--node-left p2) br)
 261         (setf (avl-tree--node-balance br) (if (> b2 0) -1 0))
 262         (setf (avl-tree--node-balance p1) (if (< b2 0) +1 0))
 263         (setf (avl-tree--node-branch node branch) p2))
 264       (setf (avl-tree--node-balance (avl-tree--node-branch node branch)) 0)
 265       nil))))
 266
 267 (defun avl-tree--enter-balance2 (node branch)
 268   ;; Return t if the tree has grown.
 269   (let ((br (avl-tree--node-branch node branch))
 270         p1 p2 b2)
 271     (cond
 272      ((> (avl-tree--node-balance br) 0)
 273       (setf (avl-tree--node-balance br) 0)
 274       nil)
 275
 276      ((= (avl-tree--node-balance br) 0)
 277       (setf (avl-tree--node-balance br) -1)
 278       t)
 279
 280      (t
 281       ;; Balance was -1 => Rebalance.
 282       (setq p1 (avl-tree--node-left br))
 283       (if (< (avl-tree--node-balance p1) 0)
 284           ;; Single LL rotation.
 285           (progn
 286             (setf (avl-tree--node-left br) (avl-tree--node-right p1))
 287             (setf (avl-tree--node-right p1) br)
 288             (setf (avl-tree--node-balance br) 0)
 289             (setf (avl-tree--node-branch node branch) p1))
 290
 291         ;; Double LR rotation.
 292         (setq p2 (avl-tree--node-right p1)
 293               b2 (avl-tree--node-balance p2))
 294         (setf (avl-tree--node-right p1) (avl-tree--node-left p2))
 295         (setf (avl-tree--node-left p2) p1)
 296         (setf (avl-tree--node-left br) (avl-tree--node-right p2))
 297         (setf (avl-tree--node-right p2) br)
 298         (setf (avl-tree--node-balance br) (if (< b2 0) +1 0))
 299         (setf (avl-tree--node-balance p1) (if (> b2 0) -1 0))
 300         (setf (avl-tree--node-branch node branch) p2))
 301       (setf (avl-tree--node-balance (avl-tree--node-branch node branch)) 0)
 302       nil))))
 303
 304 (defun avl-tree--do-enter (cmpfun root branch data)
 305   ;; Return t if height of tree ROOT has grown.  INTERNAL USE ONLY.
 306   (let ((br (avl-tree--node-branch root branch)))
 307     (cond
 308      ((null br)
 309       ;; Data not in tree, insert it.
 310       (setf (avl-tree--node-branch root branch)
 311             (avl-tree--node-create nil nil data 0))
 312       t)
 313
 314      ((funcall cmpfun data (avl-tree--node-data br))
 315       (and (avl-tree--do-enter cmpfun br 0 data)
 316            (avl-tree--enter-balance2 root branch)))
 317
 318      ((funcall cmpfun (avl-tree--node-data br) data)
 319       (and (avl-tree--do-enter cmpfun br 1 data)
 320            (avl-tree--enter-balance1 root branch)))
 321
 322      (t
 323       (setf (avl-tree--node-data br) data)
 324       nil))))
 325
 326 ;; ----------------------------------------------------------------
 327
 328 (defun avl-tree--mapc (map-function root)
 329   ;; Apply MAP-FUNCTION to all nodes in the tree starting with ROOT.
 330   ;; The function is applied in-order.
 331   ;;
 332   ;; Note: MAP-FUNCTION is applied to the node and not to the data itself.
 333   ;; INTERNAL USE ONLY.
 334   (let ((node root)
 335         (stack nil)
 336         (go-left t))
 337     (push nil stack)
 338     (while node
 339       (if (and go-left
 340                (avl-tree--node-left node))
 341           ;; Do the left subtree first.
 342           (progn
 343             (push node stack)
 344             (setq node (avl-tree--node-left node)))
 345         ;; Apply the function...
 346         (funcall map-function node)
 347         ;; and do the right subtree.
 348         (setq node (if (setq go-left (avl-tree--node-right node))
 349                        (avl-tree--node-right node)
 350                      (pop stack)))))))
 351
 352 (defun avl-tree--do-copy (root)
 353   ;; Copy the avl tree with ROOT as root.
 354   ;; Highly recursive. INTERNAL USE ONLY.
 355   (if (null root)
 356       nil
 357     (avl-tree--node-create
 358      (avl-tree--do-copy (avl-tree--node-left root))
 359      (avl-tree--do-copy (avl-tree--node-right root))
 360      (avl-tree--node-data root)
 361      (avl-tree--node-balance root))))
 362
 363 \f
 364 ;; ================================================================
 365 ;;; The public functions which operate on AVL trees.
 366
 367 (defalias 'avl-tree-compare-function 'avl-tree--cmpfun
 368   "Return the comparison function for the avl tree TREE.
 369
 370 \(fn TREE)")
 371
 372 (defun avl-tree-empty (tree)
 373   "Return t if avl tree TREE is emtpy, otherwise return nil."
 374   (null (avl-tree--root tree)))
 375
 376 (defun avl-tree-enter (tree data)
 377   "In the avl tree TREE insert DATA.
 378 Return DATA."
 379   (avl-tree--do-enter (avl-tree--cmpfun tree)
 380                       (avl-tree--dummyroot tree)
 381                       0
 382                       data)
 383   data)
 384
 385 (defun avl-tree-delete (tree data)
 386   "From the avl tree TREE, delete DATA.
 387 Return the element in TREE which matched DATA,
 388 nil if no element matched."
 389   (avl-tree--do-delete (avl-tree--cmpfun tree)
 390                        (avl-tree--dummyroot tree)
 391                        0
 392                        data))
 393
 394 (defun avl-tree-member (tree data)
 395   "Return the element in the avl tree TREE which matches DATA.
 396 Matching uses the compare function previously specified in
 397 `avl-tree-create' when TREE was created.
 398
 399 If there is no such element in the tree, the value is nil."
 400   (let ((node (avl-tree--root tree))
 401         (compare-function (avl-tree--cmpfun tree))
 402         found)
 403     (while (and node
 404                 (not found))
 405       (cond
 406        ((funcall compare-function data (avl-tree--node-data node))
 407         (setq node (avl-tree--node-left node)))
 408        ((funcall compare-function (avl-tree--node-data node) data)
 409         (setq node (avl-tree--node-right node)))
 410        (t
 411         (setq found t))))
 412     (if node
 413         (avl-tree--node-data node)
 414       nil)))
 415
 416 (defun avl-tree-map (__map-function__ tree)
 417   "Apply __MAP-FUNCTION__ to all elements in the avl tree TREE."
 418   (avl-tree--mapc
 419    (lambda (node)
 420      (setf (avl-tree--node-data node)
 421            (funcall __map-function__ (avl-tree--node-data node))))
 422    (avl-tree--root tree)))
 423
 424 (defun avl-tree-first (tree)
 425   "Return the first element in TREE, or nil if TREE is empty."
 426   (let ((node (avl-tree--root tree)))
 427     (when node
 428       (while (avl-tree--node-left node)
 429         (setq node (avl-tree--node-left node)))
 430       (avl-tree--node-data node))))
 431
 432 (defun avl-tree-last (tree)
 433   "Return the last element in TREE, or nil if TREE is empty."
 434   (let ((node (avl-tree--root tree)))
 435     (when node
 436       (while (avl-tree--node-right node)
 437         (setq node (avl-tree--node-right node)))
 438       (avl-tree--node-data node))))
 439
 440 (defun avl-tree-copy (tree)
 441   "Return a copy of the avl tree TREE."
 442   (let ((new-tree (avl-tree-create (avl-tree--cmpfun tree))))
 443     (setf (avl-tree--root new-tree) (avl-tree--do-copy (avl-tree--root tree)))
 444     new-tree))
 445
 446 (defun avl-tree-flatten (tree)
 447   "Return a sorted list containing all elements of TREE."
 448   (nreverse
 449    (let ((treelist nil))
 450      (avl-tree--mapc
 451       (lambda (node) (push (avl-tree--node-data node) treelist))
 452       (avl-tree--root tree))
 453      treelist)))
 454
 455 (defun avl-tree-size (tree)
 456   "Return the number of elements in TREE."
 457   (let ((treesize 0))
 458     (avl-tree--mapc
 459      (lambda (data) (setq treesize (1+ treesize)))
 460      (avl-tree--root tree))
 461     treesize))
 462
 463 (defun avl-tree-clear (tree)
 464   "Clear the avl tree TREE."
 465   (setf (avl-tree--root tree) nil))
 466
 467 (provide 'avl-tree)
 468
 469 ;; arch-tag: 47e26701-43c9-4222-bd79-739eac6357a9
 470 ;;; avl-tree.el ends here