Created
April 9, 2024 12:48
-
-
Save vanaur/bea2b0ea57b58140b9c1d39a9b02e998 to your computer and use it in GitHub Desktop.
This module implements a generic union-find data structure in F# (for .NET)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
/// This module implements a generic [union-find](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) data structure. | |
/// This structure is used to manage disjoint sets, often used in partitioning or clustering problems. | |
/// | |
/// Here is a summary of the features and members of this data structure: | |
/// | |
/// 1. **Constructor**: The `UnionFind` constructor initializes a new instance of the Union-Find data structure with | |
/// default elements of the original set specified as a sequence. This set can be extended. | |
/// | |
/// 2. **Public Members** : | |
/// - `UnionFind.AddElement` Adds a new element to the set. | |
/// - `UnionFind.AddRange` Adds a sequence of elements to the set. | |
/// - `UnionFind.Connect`: Combines two equivalence classes to which two given representatives belong into a single class. | |
/// - `UnionFind.ConnectRange`: Combines all equivalence classes to which the given representatives belong into a single class. | |
/// - `UnionFind.GetConnections`: Retrieves all elements of the equivalence class to which a given element belongs. | |
/// - `UnionFind.GetAllConnections`: Retrieves all equivalence classes. | |
/// - `UnionFind.AreConnected`: Checks whether two elements belong to the same equivalence class. | |
/// | |
/// 3. **Inline functions**: These functions allow you to access members of the `UnionFind` instance more concisely. | |
/// | |
/// Example: | |
/// <example> | |
/// let uf = UnionFind [| 1..6 |] | |
/// connectRange uf [| 1; 2; 3 |] | |
/// connectRange uf [| 4; 5 |] | |
/// </example> | |
/// | |
module public UnionFind = | |
open System.Collections.Generic | |
open System.Linq | |
/// Generic [union-find](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) data structure | |
/// using the path compression find amelioration, making `Find` and `Connect` O(log n) in the worst case. | |
type public UnionFind<'t when 't: equality>(defaultElements: 't seq) = | |
let mutable parents = Dictionary<'t, 't>() | |
let mutable ranks = Dictionary<'t, int>() | |
do | |
Seq.iter | |
(fun element -> | |
parents.Add(element, element) | |
ranks.Add(element, 1)) | |
defaultElements | |
let find (element: 't) = | |
let rec inner (element: 't) = | |
match parents.TryGetValue element with | |
| true, parent when parent <> element -> | |
let grandParent = parents.[parent] | |
parents.[element] <- grandParent | |
inner grandParent | |
| _ -> element | |
inner element | |
let treeSize (element: 't) = ranks.[element] | |
member internal __.GetParents | |
with get () = parents | |
and set value = parents <- value | |
member internal __.GetRanks | |
with get () = ranks | |
and set value = ranks <- value | |
/// Adds an element to the original set. | |
member public __.AddElement(element: 't) = | |
match parents.TryGetValue element with | |
| false, _ -> | |
parents.Add(element, element) | |
ranks.Add(element, 1) | |
| _ -> () | |
/// Adds the list of elements given in the original set. | |
member public this.AddRange(elements: 't seq) = Seq.iter this.AddElement elements | |
/// Combines the two equivalence classes to which the two given representatives belong into a single class. | |
member public __.Connect (left: 't) (right: 't) = | |
assert (parents.ContainsKey left) | |
assert (parents.ContainsKey right) | |
let leftTree = find left | |
let rightTree = find right | |
if leftTree = rightTree then | |
() | |
let leftTreeSize = treeSize leftTree | |
let rightTreeSize = treeSize rightTree | |
if leftTreeSize < rightTreeSize then | |
parents.[leftTree] <- rightTree | |
ranks.[rightTree] <- leftTreeSize + rightTreeSize | |
else | |
parents.[rightTree] <- leftTree | |
ranks.[leftTree] <- leftTreeSize + rightTreeSize | |
/// Combines all the equivalence classes to which the list of given representatives belong into a single class. | |
member public this.ConnectRange(xs: 't seq) = | |
if Seq.length xs < 2 then | |
() | |
let head = Seq.head xs | |
let tail = Seq.tail xs | |
Seq.iter (this.Connect head) tail | |
/// Retrieves the equivalence class to which the given element belongs. | |
member public __.GetConnections(item: 't) = | |
if not (parents.ContainsKey item) then | |
Array.empty | |
else | |
let representative = find item | |
parents | |
.GroupBy(fun kvp -> kvp.Value) | |
.Where(fun x -> x.Key = representative) | |
.Select(fun gp -> gp.Select(fun kvp -> kvp.Key) |> Seq.toArray) | |
|> Seq.toArray | |
/// Gets the set of equivalence classes. | |
member public __.GetAllConnections() = | |
parents | |
.GroupBy(fun kvp -> kvp.Value) | |
.Select(fun gp -> gp.Select(fun kvp -> kvp.Key) |> Seq.toArray) | |
|> Seq.toArray | |
/// Check whether two elements belong to the same equivalence class. | |
member public __.AreConnected (left: 't) (right: 't) = find left = find right | |
/// Adds an element to the original set. | |
let inline public addElement (uf: 't UnionFind) = uf.AddElement | |
/// Adds the list of elements given in the original set. | |
let inline public addRange (uf: 't UnionFind) = uf.AddRange | |
/// Combines the two equivalence classes to which the two given representatives belong into a single class. | |
let inline public connect (uf: 't UnionFind) = uf.Connect | |
/// Combines all the equivalence classes to which the list of given representatives belong into a single class. | |
let inline public connectRange (uf: 't UnionFind) = uf.ConnectRange | |
/// Retrieves the equivalence class to which the given element belongs. | |
let inline public getConnections (uf: 't UnionFind) = uf.GetConnections | |
/// Gets the set of equivalence classes. | |
let inline public getAllConnections (uf: 't UnionFind) = uf.GetAllConnections() | |
/// Check whether two elements belong to the same equivalence class. | |
let inline public areConnected (uf: 't UnionFind) = uf.AreConnected |
Example with a custom data type:
[<CustomEquality; NoComparison>]
type MyType =
| A
| B
| C
| D
| E
| F
override this.Equals obj =
match obj with
| :? MyType as letter ->
match this, letter with
| A, A -> true
| B, B -> true
| C, C -> true
| D, D -> true
| E, E -> true
| F, F -> true
| _ -> false
| _ -> false
override this.GetHashCode() =
match this with
| A -> 1
| B -> 2
| C -> 3
| D -> 4
| E -> 5
| F -> 6
let uf2 = UnionFind [ A; B; C; D; E; F ]
uf2.ConnectRange [ A; B; F ]
uf2.ConnectRange [ C; E ]
printfn "Are A and E in the same equivalence class? %b" (uf2.AreConnected A E)
printfn "List of equivalences classes: %A" (uf2.GetAllConnections())
Output:
Are A and E in the same equivalence class? false
List of equivalences classes: [|[|A; B; F|]; [|C; E|]; [|D|]|]
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Example of usage with integers:
Output: