module dsets;

/// dsets is an implementation of disjoint sets. It is implemented
/// with a simple class. To construct it, you provide the maximum node
/// number.
public class DSets
{
private:
    ulong[] ids, heights;
public:
    /// A `DSets(10)` would create a disjoint set of 0 to 10
    /// inclusive.
    this(ulong maxN)
    {
        ids.reserve(maxN + 1);
        heights.reserve(maxN + 1);

        foreach (n; 0 .. maxN + 1)
        {
            ids ~= n;
            heights ~= 0;
        }
    }

    /// Returns the root id of a given node
    ulong getRoot(ulong n) const
    {
        ulong result = n;
        while (ids[result] != result)
            result = ids[result];
        return result;
    }

    /// Connect two nodes. Does not check membership so may throw
    /// errors
    void connect(ulong n, ulong m)
    {
        auto nRoot = getRoot(n), mRoot = getRoot(m);
        if (nRoot != mRoot)
        {
            if (heights[nRoot] >= heights[mRoot])
            {
                ids[mRoot] = ids[nRoot];
                heights[nRoot] += heights[mRoot];
            }
            else
            {
                ids[nRoot] = ids[mRoot];
                heights[mRoot] += heights[nRoot];
            }
        }
    }

    /// Check if two given nodes are connected.
    bool connected(ulong n, ulong m) const
    {
        return getRoot(n) == getRoot(m);
    }

    ulong length() const @safe
    {
        return ids.length;
    }

    // TODO: implement disconnectNode, toArray

    /// Add a node
    void grow(ulong nodes = 0)
    {
        auto toAdd = ids[$ - 1];
        foreach (_; 0 .. nodes + 1)
        {
            ids ~= toAdd++;
            heights ~= 0;
        }
    }

}

unittest
{
    auto forrest = new DSets(9);

    forrest.connect(9, 3);
    forrest.connect(1, 3);
    forrest.connect(9, 0);
    forrest.connect(7, 5);
    forrest.connect(6, 4);
    assert(forrest.connected(1, 0));
    assert(forrest.connected(3, 0));
    assert(!forrest.connected(5, 4));
}