DisjointSets.cs

using System;
using System.Collections.Generic;
using System.Linq;

namespace mapf
{
    public class DisjointSetItem
    {
        public DisjointSetItem parent;
        public int rank;

        public DisjointSetItem()
        {
            parent = this;
            rank = 0;
        }

        public bool IsSingle()
        {
            if (this.parent != this)
                return false;
            if (this.rank != 0) // Rank is only correct in representative items!
                return false;
            return true;
        }
    }

    /// <summary>
    /// Algorithm from http://en.wikipedia.org/wiki/Disjoint-set_data_structure.
    /// Operations cost O(Ackerman^-1(n)) time - practically constant. Cost is asymptotically optimal :)
    /// </summary>
    public class DisjointSets<T>
    {
        protected IDictionary<T, DisjointSetItem> entriesToItems;
        public int maxRank { get; protected set; }

        public DisjointSets()
        {
            entriesToItems = new Dictionary<T, DisjointSetItem>();
            maxRank = 0;
        }

        /// <summary>
        /// Initially all entries are single item sets.
        /// </summary>
        /// <param name="entries"></param>
        public DisjointSets(ISet<T> entries)
        {
            entriesToItems = new Dictionary<T, DisjointSetItem>();
            maxRank = 0;
            foreach (var entry in entries)
            {
                entriesToItems[entry] = new DisjointSetItem();
            }
        }

        public DisjointSets(DisjointSets<T> other)
        {
            entriesToItems = new Dictionary<T, DisjointSetItem>();
            maxRank = other.maxRank;
            var queue = new Queue<T>(other.entriesToItems.Keys);
            var otherRepsToNewReps = new Dictionary<DisjointSetItem, DisjointSetItem>();
            while (queue.Count != 0)
            {
                var entry = queue.Dequeue();
                var item = other.entriesToItems[entry];
                other.Find(item); // Makes item.parent point directly do the set's representative
                if (item.parent == item) // This is the set's representative
                {
                    var newItem = new DisjointSetItem();
                    entriesToItems[entry] = newItem;
                    // .parent already points to itself as needed
                    newItem.rank = item.rank; // Since all items in a set do not point directly to the set's rep,
                                              // the set's rank would have been 1 regardless of its size, had it
                                              // been created by calls to Union. We prefer to remember the original rank.
                    otherRepsToNewReps[item] = newItem;
                }
                else
                {
                    if (otherRepsToNewReps.ContainsKey(item.parent)) // We've copied this rep already
                    {
                        var newItem = new DisjointSetItem();
                        entriesToItems[entry] = newItem;
                        newItem.parent = otherRepsToNewReps[item.parent];
                        // .rank not important in non-rep items
                    }
                    else
                        queue.Enqueue(entry); // Wait for the rep to be copied first.
                }
            }
        }

        /// <summary>
        /// Does nothing if entry is already present.
        /// Returns whether the entry was really added now.
        /// </summary>
        public bool AddSet(T entry)
        {
            if (entriesToItems.ContainsKey(entry) == false)
            {
                entriesToItems[entry] = new DisjointSetItem();
                return true;
            }
            return false;
        }

        /// <summary>
        /// Assumes entry is already in the DisjointSets data structure
        /// </summary>
        /// <param name="entry"></param>
        /// <returns></returns>
        public bool IsSingle(T entry)
        {
            var item = entriesToItems[entry];
            return item.IsSingle();
        }

        public IEnumerable<ISet<T>> GetSets()
        {
            Dictionary<DisjointSetItem, ISet<T>> repsToSets = new Dictionary<DisjointSetItem, ISet<T>>();
            foreach (var entryAndSetItem in entriesToItems)
            {
                var entry = entryAndSetItem.Key;
                var item = entryAndSetItem.Value;
                var rep = Find(item);
                if (repsToSets.ContainsKey(rep) == false)
                    repsToSets.Add(rep, new HashSet<T>());
                repsToSets[rep].Add(entry);
            }
            return repsToSets.Values;
        }

        public int GetNumOfSets()
        {
            int count = 0;
            foreach (var entryAndSetItem in entriesToItems)
            {
                var entry = entryAndSetItem.Key;
                var item = entryAndSetItem.Value;
                if (item.parent == item)
                    count += 1;
            }
            return count;
        }

        public override string ToString()
        {
            string ans = "{";
            foreach (var set in GetSets())
            {
                ans += "{";
                foreach (var item in set)
                {
                    ans += item.ToString() + ", ";
                }
                ans += "}, ";
            }
            ans += "}";
            return ans;
        }

        /// <summary>
        /// Modifies this data structure to also contains all unions from other.
        /// </summary>
        /// <param name="other"></param>
        /// <returns>Whether this object was changed</returns>
        public bool CopyUnions(DisjointSets<T> other)
        {
            // Create a temporary reverse mapping of other's items to its entries - O(n)
            var otherItemsToEntries = new Dictionary<DisjointSetItem, T>();
            foreach (var kvp in other.entriesToItems)
            {
                otherItemsToEntries[kvp.Value] = kvp.Key;
            }

            // Unite each of other's entries with its representative only - O(n)
            bool ret = false;
            foreach (var entryAndSetItem in other.entriesToItems)
            {
                var entry = entryAndSetItem.Key;
                var item = entryAndSetItem.Value;
                if (item.IsSingle() == true)
                    continue;
                other.Find(item);
                var repEntry = otherItemsToEntries[item.parent];
                bool wasOnlyUnitedNow = this.Union(entry, repEntry);
                ret = ret || wasOnlyUnitedNow;
            }
            return ret;
        }

        /// <summary>
        /// Finds the representative of the item's set
        /// </summary>
        /// <param name="x"></param>
        /// <returns></returns>
        protected DisjointSetItem Find(DisjointSetItem x)
        {
            if (x.parent != x)
                x.parent = Find(x.parent);
            return x.parent;
        }

        /// <summary>
        /// Doesn't assume x, y already in the data structure.
        /// Returns whether they were really only united now.
        /// </summary>
        /// <param name="entryA"></param>
        /// <param name="entryB"></param>
        /// <returns></returns>
        public bool Union(T entryA, T entryB)
        {
            AddSet(entryA);
            AddSet(entryB);
            var x = entriesToItems[entryA];
            var y = entriesToItems[entryB];
            var xRoot = Find(x);
            var yRoot = Find(y);
            if (xRoot == yRoot)
                return false;

            // x and y are not already in same set. Merge them.
            if (xRoot.rank < yRoot.rank)
                xRoot.parent = yRoot;
            else if (xRoot.rank > yRoot.rank)
                yRoot.parent = xRoot;
            else
            {
                // Arbitrarily choose x to be the root of the union
                yRoot.parent = xRoot;
                xRoot.rank = xRoot.rank + 1;
                if (xRoot.rank > maxRank)
                    maxRank = xRoot.rank;
            }
            return true;
        }

        /// <summary>
        /// Assumes entries are already in the DisjointSets data structure
        /// </summary>
        /// <param name="entryA"></param>
        /// <param name="entryB"></param>
        /// <returns></returns>
        public bool AreUnited(T entryA, T entryB)
        {
            var x = entriesToItems[entryA];
            var y = entriesToItems[entryB];
            var xRoot = Find(x);
            var yRoot = Find(y);
            return xRoot == yRoot;
        }

        public bool Contains(T entry)
        {
            return entriesToItems.ContainsKey(entry);
        }
    }
}