DisjointSets.cs

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());
        }
    }
}