@@ -4,12 +4,14 @@ defmodule AdventOfCode.Algorithms.DisjointSet do
4
4
5
5
More on this: https://en.wikipedia.org/wiki/Disjoint_sets
6
6
"""
7
+ alias AdventOfCode.Algorithms.DisjointSet
7
8
defstruct parents: % { } , ranks: % { }
8
9
9
- @ type mapped_array ( ) :: % { required ( non_neg_integer ( ) ) => non_neg_integer ( ) }
10
- @ type value ( ) :: non_neg_integer ( )
10
+ @ type mapped_parents ( ) :: % { term ( ) => term ( ) }
11
+ @ type mapped_array ( ) :: % { required ( non_neg_integer ( ) ) => term ( ) }
12
+ @ type value ( ) :: term ( )
11
13
@ type t ( ) :: % __MODULE__ {
12
- parents: mapped_array ( ) ,
14
+ parents: mapped_parents ( ) ,
13
15
ranks: mapped_array ( )
14
16
}
15
17
@@ -28,7 +30,7 @@ defmodule AdventOfCode.Algorithms.DisjointSet do
28
30
}
29
31
30
32
"""
31
- @ spec new ( non_neg_integer ( ) | List . t ( ) ) :: t ( )
33
+ @ spec new ( non_neg_integer ( ) | [ term ( ) ] ) :: t ( )
32
34
def new ( 0 ) , do: % __MODULE__ { }
33
35
34
36
def new ( size ) when is_integer ( size ) do
@@ -104,8 +106,46 @@ defmodule AdventOfCode.Algorithms.DisjointSet do
104
106
end
105
107
106
108
@ doc """
107
- Performs a union between two elements and returns the updated set. `:error` case is matched so that it fails
108
- in a piped flow.
109
+ Performs a union between two elements and returns a tuple with set and find status.
110
+
111
+ If either of the element being unionized were not a part of the set, then it returns `:error` as first element
112
+ of the tuple otherwise `:ok`
113
+
114
+ ## Example
115
+
116
+ iex> set = DisjointSet.new(5)
117
+ iex> set =
118
+ ...> set
119
+ ...> |> DisjointSet.union(0, 2)
120
+ ...> |> DisjointSet.union(4, 2)
121
+ ...> |> DisjointSet.union(3, 1)
122
+ iex> set
123
+ %DisjointSet{
124
+ parents: %{0 => 0, 1 => 3, 2 => 0, 3 => 3, 4 => 0},
125
+ ranks: %{0 => 2, 1 => 1, 2 => 1, 3 => 2, 4 => 1}
126
+ }
127
+ iex> DisjointSet.strict_union(set, 3, 1) == {:ok, set}
128
+ true
129
+
130
+ iex> ds = DisjointSet.new(1)
131
+ iex> DisjointSet.strict_union(ds, 100, 200) == {:error, ds}
132
+ true
133
+
134
+ """
135
+ @ spec strict_union ( t ( ) , value ( ) , value ( ) ) :: { :ok , t ( ) } | { :error , t ( ) }
136
+ def strict_union ( % __MODULE__ { } = disjoint_set , a , b ) do
137
+ with { root_a , disjoint_set_1 } <- find ( disjoint_set , a ) ,
138
+ { root_b , disjoint_set_2 } <- find ( disjoint_set_1 , b ) do
139
+ { :ok , union_by_rank ( disjoint_set_2 , root_a , root_b ) }
140
+ else
141
+ _ -> { :error , disjoint_set }
142
+ end
143
+ end
144
+
145
+ @ doc """
146
+ Performs a union between two elements and returns the updated set. It returns the set, either original or
147
+ updated depending on whether there was a match or not. See `DisjointSet.strict_union` if you want to know
148
+ element membership.
109
149
110
150
## Example
111
151
@@ -130,16 +170,10 @@ defmodule AdventOfCode.Algorithms.DisjointSet do
130
170
"""
131
171
@ spec union ( t ( ) , value ( ) , value ( ) ) :: t ( )
132
172
def union ( % __MODULE__ { } = disjoint_set , a , b ) do
133
- with { root_a , disjoint_set } <- find ( disjoint_set , a ) ,
134
- { root_b , disjoint_set } <- find ( disjoint_set , b ) do
135
- union_by_rank ( disjoint_set , root_a , root_b )
136
- else
137
- _ -> disjoint_set
138
- end
173
+ { _ , disjoint_set } = DisjointSet . strict_union ( disjoint_set , a , b )
174
+ disjoint_set
139
175
end
140
176
141
- def union ( :error , _ , _ ) , do: :error
142
-
143
177
@ doc """
144
178
Returns the connected components of a set of data. `:error` case is matched so that it fails
145
179
in a piped flow.
@@ -160,8 +194,9 @@ defmodule AdventOfCode.Algorithms.DisjointSet do
160
194
161
195
"""
162
196
@ spec components ( t ( ) | :error ) :: [ [ term ( ) ] ]
163
- def components ( % __MODULE__ { parents: parents } ) do
197
+ def components ( % __MODULE__ { parents: parents } = disjoint_set ) do
164
198
parents
199
+ |> Enum . map ( fn { k , _ } -> { k , find ( disjoint_set , k ) |> elem ( 0 ) } end )
165
200
|> Enum . group_by ( & elem ( & 1 , 1 ) , fn { a , _ } -> a end )
166
201
|> Map . values ( )
167
202
|> Enum . map ( & Enum . into ( & 1 , % MapSet { } ) )
0 commit comments