How to efficiently implement Disjoint Set/Union Find data structure?

Question

The data is a list of lists (sets), where all the elements are distinct. There are basically two operations.

• Find the position that contains a certain element. (And fast accessing the set with this position is needed also)

• Union (Join) the two subsets that contain two specified elements.

For example, with a test list: {{1}, {4,5}, {2}, {3}}, the find 5 operation will return index 2 and union 2 and 5 operation will return {{1}, {4,5,2}, {3}}, and the order of the sets and the order of elements in the sets do not matter. More information can be found here.

Current implementations are of OOP style. And when I try to write similar functions, I found that copying the complete list many times seems to be a big problem. Here is a version (very slow for my real problem) demonstrating the operations.

unionFind[data_, e_] := 
  FirstPosition[data, _?(MemberQ[#, e] &), {0}, 1, Heads -> False];
unionJoin[data_, {e1_, e2_}] := 
 Module[{e1pos, e2pos}, e1pos = unionFind[data, e1]; 
  e2pos = unionFind[data, e2]; 
  Append[Delete[data, {e1pos, e2pos}], 
   Join @@ Extract[data, {e1pos, e2pos}]]]

How to implement the two functionalities with more efficient methods?

Answer 1

You may use HoldFirst to address copying the large list.

ClearAll[findJoin];
Attributes[findJoin] = {HoldFirst};
findJoin[list_, v_] :=
 Module[{vPos, fPos},
  vPos = First@FirstPosition[list, v];
  fPos = First@FirstPosition[list, vPos];
  If[vPos != fPos,
   list[[vPos]] = Join @@ list[[{vPos, fPos}]];
   list[[fPos]] = Nothing;
   list = list; (*Activate the Nothing*)
   ];
  ]

Then

u = {{1}, {4, 5}, {2}, {3}};
findJoin[u, 5];

u

(* {{1}, {4, 5, 2}, {3}} *)

I'm not certain how the list = list bit to activate the Nothing will affect the memory. I don't think there is a function to shrink a list of list by effectively deallocating a row and updating the next pointer of the prior row to the row following the deallocated row. I think that maybe too low level for Wolfram Language directly. Though I believe you can write a Symbolic C function in Mathematica and compile it for this.

Hope this helps.

Answer 2

Here is an approach where I use two associations. One uses elements as keys, and set index as the value, and the other uses set index as a key, and set members as the value. This way both set index extraction, and set member extraction are order 1 operations. The union operation will be slower, as it will depend on the sizes of the sets. These associations will be encapsulated in a DisjointSetObject, and the function DisjointSet will create the container. First, the definition of DisjointSet:

DisjointSet[set_] := With[{vars = {Unique["Index"], Unique["Set"]}},
    vars = {
        Association @ Map[Thread] @ Thread[set -> Range @ Length @ set],
        AssociationThread[Range @ Length @ set, set]
    };
    DisjointSetObject[vars, Sequence @@ vars]
]

DisjointSet will create a DisjointSetObject container from a disjoint list of sets. For example:

ds = DisjointSet[{{1,9,2}, {3, 22, 4}}]

DisjointSetObject[{Index43, Set44}, <|1 -> 1, 9 -> 1, 2 -> 1, 3 -> 2, 22 -> 2, 4 -> 2|>, <|1 -> {1, 9, 2}, 2 -> {3, 22, 4}|>]

The association Index43 returns the set index associated with a set element, while Set44 returns the members of a given set index. For example:

Index43[9]
Set44[2]

1

{3, 22, 4}

Let's add some formatting so that we don't have to see the guts of the DisjointSetObject:

SetAttributes[DisjointSetObject, HoldFirst]

MakeBoxes[obj:DisjointSetObject[{index_, sets_}, __], StandardForm] ^:= With[
    {
    above = {
        {BoxForm`SummaryItem[{"Elements: ", Length[index]}]},
        {BoxForm`SummaryItem[{"Sets: ", Length[sets]}]}}
    },

    BoxForm`ArrangeSummaryBox[
        DisjointSetObject,
        obj,
        None,
        above,
        {},
        StandardForm,
        "Interpretable"->Automatic
    ]
]

We also need to add accessor functions extract the needed information from a DisjointSetObject:

DisjointSetObject[{index_, sets_},__]["Index", i_]:=Lookup[i] @ index
DisjointSetObject[{index_, sets_},__]["Index"]:=Values[index]
DisjointSetObject[{index_, sets_},__]["Set",i_]:=Lookup[i] @ sets
DisjointSetObject[{index_, sets_},__]["Set"] := Values[sets]

Examples:

ds["Index", 1]
ds["Index"]
ds["Set", 2]
ds["Set"]

1

{1, 1, 1, 2, 2, 2}

{3, 22, 4}

{{1, 9, 2}, {3, 22, 4}}

Finally, here is a union function for joining sets:

union[a_,b_][DisjointSetObject[{index_,sets_}, __]] := Module[{min, max},
    {min, max} = Sort @ Lookup[{a,b}] @ index;
    If[min =!= max,
        AssociateTo[index, Thread[sets[max] -> min]];
        sets[min] = Join[sets[min], sets[max]];
        KeyDropFrom[sets, max];
        sets[min],

        Null
    ]
]

Example:

union[9,3][ds]

{1, 9, 2, 3, 22, 4}

ds["Set"]
ds["Index", 1]

{{1, 9, 2, 3, 22, 4}}

1

Here is an example of a much larger disjoint set:

SeedRandom[1]
ds = DisjointSet @ First @ RandomPermutation[10^5]

DisjointSetObject[Elements: 100000 Sets: 16 Data not in notebook; Store now » ]

Some timings:

Length @ ds["Index"] //AbsoluteTiming
Length /@ ds["Set"] //AbsoluteTiming
ds["Index", 100] //AbsoluteTiming
ds["Set", 12] //AbsoluteTiming

Length @ union[100, 2000][ds] //AbsoluteTiming

{0.01876, 100000}

{0.000018, {9743, 30337, 46380, 6309, 3770, 2832, 190, 21, 50, 195, 50, 57, 40, 15, 6, 5}}

{5.*10^-6, 3}

{4.*10^-6, {1896, 42668, 64491, 17865, 25054, 92934, 26782, 16270, 93890, 78509, 22904, 22143, 79565, 89543, 88522, 36050, 45929, 37141, 77845, 68168, 42935, 61483, 25101, 49390, 34633, 75935, 4324, 27176, 56173, 61185, 56621, 13551, 36160, 95703, 64402, 39427, 72330, 11235, 79608, 64703, 51041, 83940, 76202, 62819, 47313, 73564, 70721, 50365, 65714, 82230, 90602, 74114, 43462, 67170, 58859, 13864, 13623}}

{0.046046, 76717}