From 8013ccf60e12305a348743c321d083adc03930e9 Mon Sep 17 00:00:00 2001
From: Robby Findler <robby@racket-lang.org>
Date: Tue, 29 Jan 2013 17:37:34 -0600
Subject: [PATCH] improvements to union find suggested by Sam and Danny

---
 collects/data/scribblings/union-find.scrbl | 28 +++++++++++++-
 collects/data/union-find.rkt               | 43 ++++++++++++++++------
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/collects/data/scribblings/union-find.scrbl b/collects/data/scribblings/union-find.scrbl
index 3a80599714..1315681215 100644
--- a/collects/data/scribblings/union-find.scrbl
+++ b/collects/data/scribblings/union-find.scrbl
@@ -22,6 +22,8 @@ These operations are not thread-safe.
 
 Makes a new set with the canonical element @racket[c].
 
+This is a constant time operation.
+                                           
 @examples[#:eval the-eval
                  (uf-new 'whale)
                  (uf-new 'dwarf-lantern)]
@@ -32,6 +34,8 @@ Makes a new set with the canonical element @racket[c].
 
 Returns @racket[#t] if @racket[x] was created with @racket[uf-new],
         and @racket[#f] otherwise.
+        
+This is a constant time operation.
 
 @examples[#:eval the-eval
   (uf-set? (uf-new 'spiny-dogfish))
@@ -40,6 +44,8 @@ Returns @racket[#t] if @racket[x] was created with @racket[uf-new],
 
 @defproc[(uf-find [a uf-set?]) any/c]{
   Returns the canonical element of @racket[a].
+
+  This is an amortized (essentially) constant time operation.
                                    
   @examples[#:eval the-eval
                    (uf-find (uf-new 'tasselled-wobbegong))]
@@ -52,6 +58,8 @@ them both have the same canonical element. Either
 of @racket[a] or @racket[b]'s  canonical elements may
 become the canonical element for the union.
 
+This is an amortized (essentially) constant time operation.
+
 @examples[#:eval the-eval
   (define a (uf-new 'sand-devil))
   (define b (uf-new 'pigeye))
@@ -61,9 +69,27 @@ become the canonical element for the union.
 ]
 }
 
+@defproc[(uf-same-set? [a uf-set?] [b uf-set?]) boolean?]{
+  Returns @racket[#t] if the sets @racket[a] and @racket[b]
+  have been unioned.
+
+This is an amortized (essentially) constant time operation.
+
+@examples[#:eval the-eval
+  (define a (uf-new 'finetooth))
+  (define b (uf-new 'speartooth))
+  (uf-same-set? a b)
+  (uf-union! a b)
+  (uf-same-set? a b)
+]
+
+  
+}
 
 @defproc[(uf-set-canonical! [a uf-set?] [c any/c]) void?]{
-  Changes @racket[a] to have a new canonical element
+  Changes @racket[a] to have a new canonical element.
+          
+This is an amortized (essentially) constant time operation.
 
  @examples[#:eval the-eval
   (define a (uf-new 'sand-devil))
diff --git a/collects/data/union-find.rkt b/collects/data/union-find.rkt
index f211de25a8..832608572e 100644
--- a/collects/data/union-find.rkt
+++ b/collects/data/union-find.rkt
@@ -1,10 +1,12 @@
 #lang racket/base
-
-(provide uf-set? 
-         uf-new
-         uf-union!
-         uf-find
-         uf-set-canonical!)
+(require racket/contract)
+(provide uf-set? uf-new)
+(provide
+ (contract-out
+  [uf-union! (-> uf-set? uf-set? void?)]
+  [uf-find (-> uf-set? any/c)]
+  [uf-set-canonical! (-> uf-set? any/c void?)]
+  [uf-same-set? (-> uf-set? uf-set? boolean?)]))
 
 (struct uf-set (x rank) #:mutable
   #:methods gen:custom-write
@@ -28,22 +30,22 @@
     [else
      (set-uf-set-x! b a)
      (when (= a-rank b-rank)
-       (set-uf-set-rank! a 1))]))
+       (set-uf-set-rank! a (+ a-rank 1)))]))
 (define (uf-find a)
   (define bx (uf-get-box a))
   (unbox bx))
 (define (uf-set-canonical! a b)
   (set-box! (uf-get-box a) b))
 (define (uf-get-box a)
-  (let loop ([a a])
+  (let loop ([a (uf-set-x a)])
     (cond
-      [(box? (uf-set-x a))
-       (uf-set-x a)]
+      [(box? a) a]
       [else
        (define fnd (loop (uf-set-x a)))
        (set-uf-set-x! a fnd)
        fnd])))
-
+(define (uf-same-set? a b)
+  (eq? (uf-get-box a) (uf-get-box b)))
 
 (module+ test
   (require rackunit
@@ -72,6 +74,13 @@
                   (uf-find b)
                   (uf-find b))
                 1)
+  (check-equal? (uf-same-set? (uf-new 1) (uf-new 2)) #f)
+  (check-equal? (uf-same-set? (uf-new 1) (uf-new 1)) #f)
+  (check-equal? (let ([a (uf-new 1)]
+                      [b (uf-new 1)])
+                  (uf-union! a b) 
+                  (uf-same-set? a b))
+                #t)
   (check-equal? (let ([sp (open-output-string)])
                   (display (uf-new "x") sp)
                   (get-output-string sp))
@@ -98,6 +107,18 @@
                   (get-output-string sp))
                 "#0=#<uf-set: #0#>")
   
+  (check-equal? (let ([a (uf-new 1)]
+                      [b (uf-new 2)]
+                      [c (uf-new 3)]
+                      [d (uf-new 4)])
+                  (uf-union! a b)
+                  (uf-union! c d)
+                  (uf-union! a c)
+                  (max (uf-set-rank a)
+                       (uf-set-rank b)
+                       (uf-set-rank c)
+                       (uf-set-rank d)))
+                2)
   
   (define (check-ranks uf)
     (let loop ([uf/box uf]