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