Skip to content

Instantly share code, notes, and snippets.

@VisualMelon VisualMelon/DisjointSets.cs Secret
Last active Jun 26, 2018

Embed
What would you like to do?
using System.Collections.Generic;
using System.Linq;
namespace Scratch
{
// not properly tested
public class DisjointSets<T>
{
private readonly Dictionary<T, int> Index = new Dictionary<T, int>();
private readonly List<T> ReverseIndex = new List<T>();
private readonly List<int> Ranks = new List<int>();
private readonly List<int> Parents = new List<int>();
public int Count => Parents.Count;
public void Add(T elem)
{
if (!Index.TryGetValue(elem, out int idx))
{
idx = Index.Count;
Index.Add(elem, idx);
ReverseIndex.Add(elem);
Parents.Add(idx);
Ranks.Add(1);
}
}
public void Join(T a, T b)
{
// union by rank
int ia = Flatten(Index[a]);
int ib = Flatten(Index[b]);
if (ia == ib)
return;
int ra = Ranks[ia];
int rb = Ranks[ib];
if (ra == rb)
{
Parents[ib] = ia;
Ranks[ia] = rb + 1;
}
else if (ra < rb)
{
Parents[ib] = ia;
}
else // rb > ra
{
Parents[ia] = ib;
}
}
private int Flatten(int e)
{
int parent = Parents[e];
if (parent == e)
{
return e;
}
else
{
e = Flatten(parent);
Parents[e] = e;
return e;
}
}
public IEnumerable<IEnumerable<T>> EnumerateSets()
{
Dictionary<int, IList<int>> sets = new Dictionary<int, IList<int>>();
for (int i = 0; i < Count; i++)
{
int root = Flatten(i);
if (sets.ContainsKey(root))
{
sets[root].Add(i);
}
else
{
sets.Add(root, new List<int>() { i });
}
}
return sets.Values.Select(s => s.Select(e => ReverseIndex[e]));
}
public static IEnumerable<IReadOnlyCollection<T>> Cluster(IEnumerable<IEnumerable<T>> groups)
{
DisjointSets<T> djs = new DisjointSets<T>();
foreach (var group in groups)
{
// take the first element in each group as the 'token' element, and join all the others with it (this could be done much faster if we allowed ourselves access to internal representation)
T tokenElement = default;
bool gotToken = false;
foreach (var elem in group)
{
djs.Add(elem);
if (gotToken)
{ // join us to the tokenElement
djs.Join(tokenElement, elem);
}
else
{ // we are the token element
tokenElement = elem;
gotToken = true;
}
}
}
return djs.EnumerateSets().Select(s => s.ToArray());
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.