public
Created

WeightedUnionFind rewritten in clojure

  • Download Gist
WeightedUnionFind.clj
Clojure
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53
(defrecord WeightedUnionFind [ids sizes])
 
(defn WeightedUnionFind-with-length [n]
(let [ids (vec (range n))
sizes (vec (repeat n 1))]
(->WeightedUnionFind ids sizes)))
 
; is there a better way to make the below methods so that they don't need
; to have the uf structure passed in each time?
 
; finds the root of p
(defn find-root [uf p]
(loop [p p]
(if (= p ((:ids uf) p))
p
(recur ((:ids uf) p)))))
 
; is p connected to q?
(defn connected? [uf p q]
(= (find-root uf p) (find-root uf q)))
 
; connect p and q, returning the result
(defn union [uf p q]
(let [i (find-root uf p) j (find-root uf q)]
(if (= i j)
uf
(if (< ((:sizes uf) i) ((:sizes uf) j))
(assoc uf :ids (assoc (:ids uf) i j) ; is this the best way to do this giant assoc?
:sizes (assoc (:sizes uf)
j (+ ((:sizes uf) i) ((:sizes uf) j))))
(assoc uf :ids (assoc (:ids uf) j i)
:sizes (assoc (:sizes uf)
i (+ ((:sizes uf) i) ((:sizes uf) j))))))))
 
; test script. sure would be nice to make this cleaner so it's not so noisy.
; going to look at the clojure.test stuff next...
(def uf (WeightedUnionFind-with-length 10)) ; test structure
(= 2 (find-root 2 uf)) ; 2 should be 2's root
(= false (connected? 1 2 uf)) ; 1 should not be connected to 2
(def uf (union 1 2 uf)) ; connect 1 and 2
(= 2 ((:sizes uf) 1)) ; 1 should be size 2
(= 1 (find-root 2 uf)) ; 1 should be 2's root
(= true (connected? 1 2 uf)) ; 1 should be connected to 2
(= false (connected? 1 3 uf)) ; 1 should not be connected to 3
(= false (connected? 1 9 uf)) ; 1 should not be connected to 9
(= 9 (find-root 9 uf)) ; 9 should be its own root
(def uf (union 1 3 uf)) ; connect 1 and 3
(def uf (union 3 9 uf)) ; connect 3 and 9
(= 4 ((:sizes uf) 1)) ; 1 should be size 4
(= 1 (find-root 9 uf)) ; 1 should be 9's root
(= true (connected? 1 9 uf)) ; 1 and 9 should be connected
(= false (connected? 1 8 uf)) ; 1 and 8 should not be connected
(print uf)

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.