WeightedTreeDisjointSet.cs

namespace AlgorithmsAndDataStructures.DataStructures.DisjointSet
{
    public class WeightedTreeDisjointSet : IDisjointSet
    {
        private readonly int[] set;
        private readonly int[] weight;

        public WeightedTreeDisjointSet(int size)
        {
            set = new int[size];
            weight = new int[size];

            for (var i = 0; i < size; i++)
            {
                set[i] = i;
            }

            for (var i = 0; i < size; i++)
            {
                weight[i] = 1;
            }
        }

        public bool Connected(int a, int b)
        {
            return FindRoot(a) == FindRoot(b);
        }

        public void Union(int a, int b)
        {
            var aRoot = FindRoot(a);
            var bRoot = FindRoot(b);

            if (weight[aRoot] >= weight[bRoot])
            {
                set[bRoot] = aRoot;
                weight[aRoot] = weight[aRoot] + weight[bRoot];
            }
            else
            {
                set[aRoot] = bRoot;
                weight[bRoot] = weight[aRoot] + weight[bRoot];
            }
        }

        private int FindRoot(int a)
        {
            var current = a;

            while (current != set[current])
            {
                current = set[set[current]];
            }

            return current;
        }
    }
}