(module union-find mzscheme
  #| Implementation of traditional union-find algorithm using mutation.
     Straightforward implementation from Cormen, Leiserson, Rivest, Stein.
     "Introduction to Algorithms" |#
  (require (lib "contract.ss")
           (lib "etc.ss"))
  ;; A forest contains a collection of its nodes keyed by element.
  ;; The elements are compared by eq?
  (define-struct forest (ht))
  ;; A node is an element, a parent node, and a numeric rank.
  (define-struct node (elt p rank))
  ;; -make-forest: -> forest
  ;; -make-forest: (union 'equal 'weak)* -> forest
  ;; Builds a new, empty forest.
  (define -make-forest
       (make-forest (make-hash-table))]
       (make-forest (apply make-hash-table flags))]))
  ;; lookup-node: forest X -> node
  ;; Returns the node that's associated with this element.
  (define (lookup-node a-forest an-elt)
    (hash-table-get (forest-ht a-forest) an-elt))
  ;; make-set: forest X -> void
  ;; Adds a new set into the forest.
  (define (make-set a-forest an-elt)
      [(hash-table-get (forest-ht a-forest) an-elt #f)
       (local ((define a-node (make-node an-elt #f 0)))
         (set-node-p! a-node a-node)
         (hash-table-put! (forest-ht a-forest) an-elt a-node))]))
  ;; find-set: forest X -> X
  ;; Returns the representative element of elt.
  (define (find-set a-forest an-elt)
    (local ((define a-node (lookup-node a-forest an-elt)))
      (node-elt (get-representative-node a-node))))
  ;; get-representative-node: forest node -> node
  ;; Returns the representative node of a-node, doing path
  ;; compression if we have to follow links.
  (define (get-representative-node a-node)
    (local ((define p (node-p a-node)))
      (cond [(eq? a-node p)
             (let ([rep (get-representative-node p)])
               ;; Path compression is here:
               (set-node-p! a-node rep)
  ;; union-set: forest X X -> void
  ;; Joins the two elements into the same set.
  (define (union-set a-forest elt1 elt2)
    (local ((define rep1 (get-representative-node
                          (lookup-node a-forest elt1)))
            (define rep2 (get-representative-node
                          (lookup-node a-forest elt2))))
        [(< (node-rank rep1) (node-rank rep2))
         (set-node-p! rep1 rep2)]
        [(> (node-rank rep1) (node-rank rep2))
         (set-node-p! rep2 rep1)]
         (set-node-p! rep1 rep2)
         (set-node-rank! rep1 (add1 (node-rank rep1)))])))
   [forest? (any/c . -> . boolean?)]
   [rename -make-forest make-forest
            (-> forest?)
            (() (listof (symbols 'equal 'weak)) . ->* . (forest?)))]
   [make-set (forest? any/c . -> . any)]
   [find-set (forest? any/c . -> . any)]
   [union-set (forest? any/c any/c . -> . void)]))