Slight rewrite and breaking change

disjoint-sets.lisp

@@ -11,14 +11,9 @@ Said otherwise: a set is identified by the id of its root set.
 
 |#
 
-(cl:in-package #:cl-user)
+(cl:in-package #:disjoint-sets)
 
-(defpackage #:disjoint-sets
-  (:use #:cl))
-
-(in-package #:disjoint-sets)
-
-(defun make-disjoint-sets (&optional number-of-sets)
+(defun make-disjoint-sets (&optional (arity 0))
   "Create a set of sets represented as an array.
 
 Examples:
@@ -28,27 +23,20 @@ Examples:
 (make-disjoint-sets 10)
 ;; => #(0 1 2 3 4 5 6 7 8 9)
 "
-  (let ((sets (make-array (list (or number-of-sets 0))
-                          :element-type 'integer
-                          :adjustable t
-                          :fill-pointer t)))
-    (when number-of-sets
-      (loop :for i :below number-of-sets
-            :do (setf (aref sets i) i)))
-    sets))
+  (let ((sets (make-array `(,arity) :element-type 'integer
+                          :adjustable t :fill-pointer t)))
+    (dotimes (i arity sets) (setf (aref sets i) i))))
 
 (defun disjoint-sets-add (sets)
-  "Add a new item into its own disjoint set. Return a new id.
+  "Add a new item into its own disjoint set. Return the new id.
 
 Example:
 
 (let ((id (disjoint-sets-add sets)))
   ;; SETS is modified
   ...)
 "
-  (let ((new-id  (length sets)))
-    (vector-push-extend new-id sets)
-    new-id))
+  (vector-push-extend (length sets) sets))
 
 (defun disjoint-sets-find (sets id)
   "Find the id of the set representative (the root).
@@ -58,28 +46,40 @@ Example:
 (disjoint-sets-find sets 5)
 "
   (let ((parent (aref sets id)))
-    (if (= id parent)
-        ;; If "id" is the root, just return it.
-        id
+    (if (= id parent) id             ; If `id' is the root, return it.
         (let ((root (disjoint-sets-find sets parent)))
           ;; Path compression: point directly to the root if it's not
           ;; already the case.
           (when (/= root parent)
             (setf (aref sets id) root))
           root))))
 
-(defun disjoint-sets-union (sets id1 id2)
+(defun disjoint-sets-unify (sets id1 id2)
   "Merge two disjoint sets. Return the set representative (the root)
 
 Example:
 
-(disjoint-sets-union sets 1 2)
-=> 4 ; SETS is modified.
+(disjoint-sets-unify sets 1 2)
+=> 4 ; `sets' is modified.
 "
   (let ((root1 (disjoint-sets-find sets id1))
         (root2 (disjoint-sets-find sets id2)))
     (setf (aref sets root2) root1)))
 
+(defun copy-disjoint-sets (sets)
+  "Copy a set of sets into a fresh array.
+
+Exmaple:
+(let* ((sets (make-disjoint-sets 2))
+       (copy (copy-disjoint-sets sets)))
+  (disjoint-sets-unify sets 0 1)
+  (values sets copy))
+=> #(0 0), #(0 1)
+"
+  (make-array (length sets) :element-type 'integer
+                            :adjustable t :fill-pointer t
+                            :initial-contents sets))
+
 (defun disjoint-sets-same-set-p (sets id1 id2)
   "Test if 2 items are in the same set.
 

package.lisp

@@ -6,5 +6,5 @@
    #:make-disjoint-sets
    #:disjoint-sets-add
    #:disjoint-sets-find
-   #:disjoint-sets-union
+   #:disjoint-sets-unify
    #:disjoint-sets-same-set-p))

tests.lisp

@@ -31,20 +31,20 @@
         (list (disjoint-sets-add sets)
               sets))))
 
-(define-test disjoint-set-union
+(define-test disjoint-set-unify
   (is equalp
       #(0 1 1 3 4 5 6 7 8 1)
       (let ((sets (make-disjoint-sets 10)))
-        (disjoint-sets-union sets 1 2)
-        (disjoint-sets-union sets 1 9)
+        (disjoint-sets-unify sets 1 2)
+        (disjoint-sets-unify sets 1 9)
         sets)))
 
 (define-test disjoint-set-same-set-p
   (is equalp
       '(t t nil)
       (let ((sets (make-disjoint-sets 10)))
-        (disjoint-sets-union sets 1 2)
-        (disjoint-sets-union sets 1 9)
+        (disjoint-sets-unify sets 1 2)
+        (disjoint-sets-unify sets 1 9)
         (list (disjoint-sets-same-set-p sets 1 2)
               (disjoint-sets-same-set-p sets 2 9)
               (disjoint-sets-same-set-p sets 1 5)))))