-
Notifications
You must be signed in to change notification settings - Fork 242
/
disjoint_set.jl
230 lines (191 loc) · 6.75 KB
/
disjoint_set.jl
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
# Disjoint-Set
############################################################
#
# A forest of disjoint sets of integers
#
# Since each element is an integer, we can use arrays
# instead of dictionary (for efficiency)
#
# Disjoint sets over other key types can be implemented
# based on an IntDisjointSet through a map from the key
# to an integer index
#
############################################################
_intdisjointset_bounds_err_msg(T) = "the maximum number of elements in IntDisjointSet{$T} is $(typemax(T))"
"""
IntDisjointSet{T<:Integer}(n::Integer)
A forest of disjoint sets of integers, which is a data structure
(also called a union–find data structure or merge–find set)
that tracks a set of elements partitioned
into a number of disjoint (non-overlapping) subsets.
"""
mutable struct IntDisjointSet{T<:Integer}
parents::Vector{T}
ranks::Vector{T}
ngroups::T
end
IntDisjointSet(n::T) where {T<:Integer} = IntDisjointSet{T}(collect(Base.OneTo(n)), zeros(T, n), n)
IntDisjointSet{T}(n::Integer) where {T<:Integer} = IntDisjointSet{T}(collect(Base.OneTo(T(n))), zeros(T, T(n)), T(n))
Base.length(s::IntDisjointSet) = length(s.parents)
"""
num_groups(s::IntDisjointSet)
Get a number of groups.
"""
num_groups(s::IntDisjointSet) = s.ngroups
Base.eltype(::Type{IntDisjointSet{T}}) where {T<:Integer} = T
# find the root element of the subset that contains x
# path compression is implemented here
function find_root_impl!(parents::Vector{T}, x::Integer) where {T<:Integer}
p = parents[x]
@inbounds if parents[p] != p
parents[x] = p = _find_root_impl!(parents, p)
end
return p
end
# unsafe version of the above
function _find_root_impl!(parents::Vector{T}, x::Integer) where {T<:Integer}
@inbounds p = parents[x]
@inbounds if parents[p] != p
parents[x] = p = _find_root_impl!(parents, p)
end
return p
end
"""
find_root!(s::IntDisjointSet{T}, x::T)
Find the root element of the subset that contains an member `x`.
Path compression happens here.
"""
find_root!(s::IntDisjointSet{T}, x::T) where {T<:Integer} = find_root_impl!(s.parents, x)
"""
in_same_set(s::IntDisjointSet{T}, x::T, y::T)
Returns `true` if `x` and `y` belong to the same subset in `s`, and `false` otherwise.
"""
in_same_set(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer} = find_root!(s, x) == find_root!(s, y)
"""
union!(s::IntDisjointSet{T}, x::T, y::T)
Merge the subset containing `x` and that containing `y` into one
and return the root of the new set.
"""
function Base.union!(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer}
parents = s.parents
xroot = find_root_impl!(parents, x)
yroot = find_root_impl!(parents, y)
return xroot != yroot ? root_union!(s, xroot, yroot) : xroot
end
"""
root_union!(s::IntDisjointSet{T}, x::T, y::T)
Form a new set that is the union of the two sets whose root elements are
`x` and `y` and return the root of the new set.
Assume `x ≠ y` (unsafe).
"""
function root_union!(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer}
parents = s.parents
rks = s.ranks
@inbounds xrank = rks[x]
@inbounds yrank = rks[y]
if xrank < yrank
x, y = y, x
elseif xrank == yrank
rks[x] += one(T)
end
@inbounds parents[y] = x
s.ngroups -= one(T)
return x
end
"""
push!(s::IntDisjointSet{T})
Make a new subset with an automatically chosen new element `x`.
Returns the new element. Throw an `ArgumentError` if the
capacity of the set would be exceeded.
"""
function Base.push!(s::IntDisjointSet{T}) where {T<:Integer}
l = length(s)
l < typemax(T) || throw(ArgumentError(_intdisjointset_bounds_err_msg(T)))
x = l + one(T)
push!(s.parents, x)
push!(s.ranks, zero(T))
s.ngroups += one(T)
return x
end
"""
DisjointSet{T}(xs)
A forest of disjoint sets of arbitrary value type `T`.
It is a wrapper of `IntDisjointSet{Int}`, which uses a
dictionary to map the input value to an internal index.
"""
mutable struct DisjointSet{T} <: AbstractSet{T}
intmap::Dict{T,Int}
revmap::Vector{T}
internal::IntDisjointSet{Int}
DisjointSet{T}() where T = new{T}(Dict{T,Int}(), Vector{T}(), IntDisjointSet(0))
function DisjointSet{T}(xs) where T # xs must be iterable
imap = Dict{T,Int}()
rmap = Vector{T}()
n = length(xs)::Int
sizehint!(imap, n)
sizehint!(rmap, n)
id = 0
for x in xs
imap[x] = (id += 1)
push!(rmap,x)
end
return new{T}(imap, rmap, IntDisjointSet(n))
end
end
DisjointSet() = DisjointSet{Any}()
DisjointSet(xs) = _DisjointSet(xs, Base.IteratorEltype(xs))
_DisjointSet(xs, ::Base.HasEltype) = DisjointSet{eltype(xs)}(xs)
function _DisjointSet(xs, ::Base.EltypeUnknown)
T = Base.@default_eltype(xs)
(isconcretetype(T) || T === Union{}) || return Base.grow_to!(DisjointSet{T}(), xs)
return DisjointSet{T}(xs)
end
Base.iterate(s::DisjointSet) = iterate(s.revmap)
Base.iterate(s::DisjointSet, i) = iterate(s.revmap, i)
Base.length(s::DisjointSet) = length(s.internal)
"""
num_groups(s::DisjointSet)
Get a number of groups.
"""
num_groups(s::DisjointSet) = num_groups(s.internal)
Base.eltype(::Type{DisjointSet{T}}) where T = T
Base.empty(s::DisjointSet{T}, ::Type{U}=T) where {T,U} = DisjointSet{U}()
function Base.sizehint!(s::DisjointSet, n::Integer)
sizehint!(s.intmap, n)
sizehint!(s.revmap, n)
return s
end
"""
find_root!{T}(s::DisjointSet{T}, x::T)
Find the root element of the subset in `s` which has the element `x` as a member.
"""
find_root!(s::DisjointSet{T}, x::T) where {T} = s.revmap[find_root!(s.internal, s.intmap[x])]
"""
in_same_set(s::DisjointSet{T}, x::T, y::T)
Return `true` if `x` and `y` belong to the same subset in `s`, and `false` otherwise.
"""
in_same_set(s::DisjointSet{T}, x::T, y::T) where {T} = in_same_set(s.internal, s.intmap[x], s.intmap[y])
"""
union!(s::DisjointSet{T}, x::T, y::T)
Merge the subset containing `x` and that containing `y` into one
and return the root of the new set.
"""
Base.union!(s::DisjointSet{T}, x::T, y::T) where {T} = s.revmap[union!(s.internal, s.intmap[x], s.intmap[y])]
"""
root_union!(s::DisjointSet{T}, x::T, y::T)
Form a new set that is the union of the two sets whose root elements are
`x` and `y` and return the root of the new set.
Assume `x ≠ y` (unsafe).
"""
root_union!(s::DisjointSet{T}, x::T, y::T) where {T} = s.revmap[root_union!(s.internal, s.intmap[x], s.intmap[y])]
"""
push!(s::DisjointSet{T}, x::T)
Make a new subset containing `x` if any existing subset of `s` does not contain `x`.
"""
function Base.push!(s::DisjointSet{T}, x::T) where T
haskey(s.intmap, x) && return x
id = push!(s.internal)
s.intmap[x] = id
push!(s.revmap, x) # Note, this assumes invariant: length(s.revmap) == id
return x
end