Create a gist now

Instantly share code, notes, and snippets.

WeightedUnionFind rewritten in clojure
(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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment