forked from dyoo/whalesong
-
Notifications
You must be signed in to change notification settings - Fork 0
/
union-find.rkt
91 lines (68 loc) · 2.43 KB
/
union-find.rkt
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#lang typed/racket/base
;; Union-find hardcoded to do symbols.
(provide (all-defined-out))
;; A forest contains a collection of its nodes keyed by element.
;; The elements are compared by eq?
(define-struct: forest
([ht : (HashTable Symbol node)]))
;; A node is an element, a parent node, and a numeric rank.
(define-struct: node
([elt : Symbol]
[p : (U False node)]
[rank : Natural])
#:mutable)
;; Builds a new, empty forest.
(: new-forest (-> forest))
(define (new-forest)
(make-forest (make-hash)))
;; lookup-node: forest X -> node
;; Returns the node that's associated with this element.
(: lookup-node (forest Symbol -> node))
(define (lookup-node a-forest an-elt)
(unless (hash-has-key? (forest-ht a-forest) an-elt)
(make-set a-forest an-elt))
(hash-ref (forest-ht a-forest)
an-elt))
;; make-set: forest X -> void
;; Adds a new set into the forest.
(: make-set (forest Symbol -> Void))
(define (make-set a-forest an-elt)
(unless (hash-has-key? (forest-ht a-forest) an-elt)
(let ([a-node (make-node an-elt #f 0)])
(set-node-p! a-node a-node)
(hash-set! (forest-ht a-forest) an-elt a-node))))
(: find-set (forest Symbol -> Symbol))
;; Returns the representative element of elt.
(define (find-set a-forest an-elt)
(let ([a-node (lookup-node a-forest an-elt)])
(node-elt (get-representative-node a-node))))
(: get-representative-node (node -> node))
;; Returns the representative node of a-node, doing path
;; compression if we have to follow links.
(define (get-representative-node a-node)
(let ([p (node-p a-node)])
(cond [(eq? a-node p)
a-node]
[(node? p)
(let ([rep (get-representative-node p)])
;; Path compression is here:
(set-node-p! a-node rep)
rep)]
[else
;; impossible situation
(error 'get-representative-node)])))
(: union-set (forest Symbol Symbol -> Void))
;; Joins the two elements into the same set.
(define (union-set a-forest elt1 elt2)
(let ([rep1 (get-representative-node
(lookup-node a-forest elt1))]
[rep2 (get-representative-node
(lookup-node a-forest elt2))])
(cond
[(< (node-rank rep1) (node-rank rep2))
(set-node-p! rep1 rep2)]
[(> (node-rank rep1) (node-rank rep2))
(set-node-p! rep2 rep1)]
[else
(set-node-p! rep1 rep2)
(set-node-rank! rep1 (add1 (node-rank rep1)))])))