Disjoint-set data structure (C++)

Question

I'm an intermediate-level C++ programmer aiming to write reference-quality code. I hope to learn about any mistakes in this piece of code and blind spots in my understanding of concepts. Notes:

• The stylistic use of public:/private: on every member is subject of a question on Software Engineering.

• I'm not totally sure if I made the best choice by having the assignment operator take an object by value, instead of defining two separate assignment operators for copy-by-reference and move-by-rvalue-reference.

• I intend to avoid all undefined behavior on all platforms. I care very much about things like integer bit width guarantees, integer type conversion rules, signed integer overflows, etc.

• If a feature like struct/class has a C flavor and C++ flavor, I generally prefer the C++ flavor. (An exception is that I prefer value construction with = instead of () or {}.)

---

Sample usage:

DisjointSet<std::uint16_t> ds(4);
ds.mergeSets(2, 0);
std::cout << ds.getNumberOfSets() << std::endl;  // 3
std::cout << ds.areInSameSet(0, 2)) << std::endl;  // true

(For further examples, see the accompanying test program.)

Library:

_download

#pragma once

#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>


/* 
 * Represents a set of disjoint sets. Also known as the union-find data structure.
 * Main operations are querying if two elements are in the same set, and merging two sets together.
 * Useful for testing graph connectivity, and is used in Kruskal's algorithm.
 * The parameter S can be any integer type, such as size_t. For any given S, the maximum number
 * of sets is S_MAX. Using a smaller type like int8_t can help save memory compared to uint64_t.
 */
template <typename S>
class DisjointSet final {
    
    /*---- Helper structure ----*/
    
    private: struct Node final {
        // The index of the parent element. An element is a representative
        // iff its parent is itself. Mutable due to path compression.
        mutable S parent;
        
        // Always in the range [0, floor(log2(numElems))]. For practical computers, this has a maximum value of 64.
        // Note that signed char is guaranteed to cover at least the range [0, 127].
        signed char rank;
        
        // Positive number if the element is a representative, otherwise zero.
        S size;
    };
    
    
    
    /*---- Fields ----*/
    
    private: std::vector<Node> nodes;
    private: S numSets;
    
    
    
    /*---- Constructors ----*/
    
    // Constructs a new set containing the given number of singleton sets.
    // For example, DisjointSet(3) --> {{0}, {1}, {2}}.
    // Even if S has a wider range than size_t, it is required that 1 <= numElems <= SIZE_MAX.
    public: explicit DisjointSet(S numElems) :
            numSets(numElems) {
        if (numElems < 0)
            throw std::domain_error("Number of elements must be non-negative");
        if (!safeLessEquals(numElems, SIZE_MAX))
            throw std::length_error("Number of elements too large");
        nodes.reserve(static_cast<std::size_t>(numElems));
        for (S i = 0; i < numElems; i++)
            nodes.push_back(Node{i, 0, 1});
    }
    
    
    public: explicit DisjointSet(const DisjointSet &other) = default;
    
    
    public: DisjointSet(DisjointSet &&other) = default;
    
    
    public: DisjointSet &operator=(DisjointSet other) {
        std::swap(nodes  , other.nodes  );
        std::swap(numSets, other.numSets);
        return *this;
    }
    
    
    
    /*---- Methods ----*/
    
    // Returns the number of elements among the set of disjoint sets; this was the number passed
    // into the constructor and is constant for the lifetime of the object. All the other methods
    // require the argument elemIndex to satisfy 0 <= elemIndex < getNumberOfElements().
    public: S getNumberOfElements() const {
        return static_cast<S>(nodes.size());
    }
    
    
    // Returns the number of disjoint sets overall. This number decreases monotonically as time progresses;
    // each call to mergeSets() either decrements the number by one or leaves it unchanged. 0 <= result <= getNumberOfElements().
    public: S getNumberOfSets() const {
        return numSets;
    }
    
    
    // Returns the size of the set that the given element is a member of. 1 <= result <= getNumberOfElements().
    public: S getSizeOfSet(S elemIndex) const {
        return nodes.at(getRepr(elemIndex)).size;
    }
    
    
    // Tests whether the given two elements are members of the same set. Note that the arguments are orderless.
    public: bool areInSameSet(S elemIndex0, S elemIndex1) const {
        return getRepr(elemIndex0) == getRepr(elemIndex1);
    }
    
    
    // Merges together the sets that the given two elements belong to. This method is also known as "union" in the literature.
    // If the two elements belong to different sets, then the two sets are merged and the method returns true.
    // Otherwise they belong in the same set, nothing is changed and the method returns false. Note that the arguments are orderless.
    public: bool mergeSets(S elemIndex0, S elemIndex1) {
        // Get representatives
        std::size_t repr0 = getRepr(elemIndex0);
        std::size_t repr1 = getRepr(elemIndex1);
        if (repr0 == repr1)
            return false;
        
        // Compare ranks
        int cmp = nodes.at(repr0).rank - nodes.at(repr1).rank;
        // Note: The computation of cmp does not overflow. 0 <= ranks[i] <= SCHAR_MAX,
        // so SCHAR_MIN <= -SCHAR_MAX <= ranks[i] - ranks[j] <= SCHAR_MAX.
        // The result actually fits in a signed char, and with sizeof(char) <= sizeof(int),
        // the promotion to int still guarantees the result fits.
        if (cmp == 0)  // Increment repr0's rank if both nodes have same rank
            nodes.at(repr0).rank++;
        else if (cmp < 0)  // Swap to ensure that repr0's rank >= repr1's rank
            std::swap(repr0, repr1);
        
        // Graft repr1's subtree onto node repr0
        nodes.at(repr1).parent = repr0;
        nodes.at(repr0).size += nodes.at(repr1).size;
        nodes.at(repr1).size = 0;
        numSets--;
        return true;
    }
    
    
    // For unit tests. This detects many but not all invalid data structures, throwing an exception
    // if a structural invariant is known to be violated. This always returns silently on a valid object.
    public: void checkStructure() const {
        S numRepr = 0;
        S i = 0;
        for (const Node &node : nodes) {
            bool isRepr = node.parent == i;
            if (isRepr)
                numRepr++;
            
            bool ok = true;
            ok &= 0 <= node.parent && safeLessThan(node.parent, nodes.size());
            ok &= 0 <= node.rank && (isRepr || node.rank < nodes.at(node.parent).rank);
            ok &= 0 <= node.size && safeLessEquals(node.size, nodes.size());
            ok &= (!isRepr && node.size == 0) || (isRepr && node.size >= (static_cast<S>(1) << node.rank));
            if (!ok)
                throw std::logic_error("Assertion error");
            i++;
        }
        if (!(0 <= numSets && numSets == numRepr && safeLessEquals(numSets, nodes.size())))
            throw std::logic_error("Assertion error");
    }
    
    
    // (Private) Returns the representative element for the set containing the given element. This method is also
    // known as "find" in the literature. Also performs path compression, which alters the internal state to
    // improve the speed of future queries, but has no externally visible effect on the values returned.
    private: S getRepr(S elemIndex) const {
        // Follow parent pointers until we reach a representative
        S parent = nodes.at(elemIndex).parent;
        while (true) {
            S grandparent = nodes.at(static_cast<std::size_t>(parent)).parent;
            if (grandparent == parent)
                return static_cast<std::size_t>(parent);
            nodes.at(static_cast<std::size_t>(elemIndex)).parent = grandparent;  // Partial path compression
            elemIndex = parent;
            parent = grandparent;
        }
    }
    
    
    private: static bool safeLessThan(S x, std::size_t y) {
        return (std::is_signed<S>::value && x < 0) ||
            static_cast<typename std::make_unsigned<S>::type>(x) < y;
    }
    
    
    private: static bool safeLessEquals(S x, std::size_t y) {
        return (std::is_signed<S>::value && x < 0) ||
            static_cast<typename std::make_unsigned<S>::type>(x) <= y;
    }
    
};

Answer 1

Here are some things that may help you improve your program.

Use include guards instead of #pragma once

The use of #pragma once, while common, is not standard C++. Use include guards instead as:

#ifndef DISJOINT_SET_HPP
#define DISJOINT_SET_HPP
// your header here
#endif // DISJOINT_SET_HPP

See SF.8 for details.

Don't mark every data member with public or private

If you are uniquely writing in this style, then almost by definition, everyone else who uses your code will find it that much more difficult to read and understand your code. Inserting unnecessary keywords adds visual clutter rather than clarity.

Avoid unneeded casts

The getRepr() function is declared as returning type S, but the return statement within it is this:

return static_cast<std::size_t>(parent);

That makes no sense at all because parent is already declared to be of type S.

Give better messages in exceptions

In the checkStructure() function, the code checks to verify a large number of predicates, collapsing them all into a single bool. If any of the conditions fails, the caller gets only the vague "Assertion error" message. At this point, the structure is somehow corrupted and no longer usable, which represents a rather catastrophic failure of the structure. For that reason, I'd suggest either writing each clause as a separate throw with a clear message or using old-fashioned assert which is likely to be the better choice in this case.

Move error checking to compile time where practical

The constructor currently looks like this:

public: explicit DisjointSet(S numElems) :
        numSets(numElems) {
    if (numElems < 0)
        throw std::domain_error("Number of elements must be non-negative");
    if (!safeLessEquals(numElems, SIZE_MAX))
        throw std::length_error("Number of elements too large");
    nodes.reserve(static_cast<std::size_t>(numElems));
    for (S i = 0; i < numElems; i++)
        nodes.push_back(Node{i, 0, 1});
}

However, if we look carefully at this, both of the checks could actually be done at compile-time instead. That suggests that the use of C++20 requires or C++11 enable_if. Here I'm using the C++14 enable_if_t which is just a convenience function:

template <typename S, 
    typename = std::enable_if_t<std::is_integral<S>::value>,
    typename = std::enable_if_t<std::is_unsigned<S>::value> 
>
class DisjointSet final {

Now if we attempt to create something invalid like DisjoinSet<float> ds(3); we get a much clearer error versus a relatively obscure and unhelpful error pointing to safeLessEquals or checkStructure.

Provide sensible default values where practical

The Node structure is only used internally to the DisjointSet class and is only created in the constructor of DisjointSet. For that reason, I'd suggest declaring it like this:

struct Node final {
    mutable S parent;
    signed char rank = 0;
    S size = 1;
    Node& operator++() { ++parent; return *this; }
};

If you're wondering about the operator, that's part of the next suggestion.

Use standard library functions where practical

I would suggest rewriting the constructor for DisjointSet like this:

explicit DisjointSet(S numElems) 
    : numSets(numElems) 
    , nodes(numElems) 
{
    std::iota(nodes.begin(), nodes.end(), Node{0});
}

Note that this also assumes that we declare numSets before nodes so as not to confuse the reader. (Elements are initialized in declaration order.) In C++20, one could likely use a range and make this even more efficient.

Don't hide a loop return condition

The getRepr() function has a while(true), but actually returns when the parent and grandparent are the same. I'd suggest rewriting that so that the looping exit condition is part of the loop rather than hidden within the body. For example, one way to rewrite would be this:

S getRepr(S elemIndex) const {
    while (nodes.at(elemIndex).parent != nodes.at(nodes.at(elemIndex).parent).parent) {
        auto parent = nodes.at(elemIndex).parent;
        // set parent = grandparent
        nodes.at(elemIndex).parent = nodes.at(parent).parent;
        elemIndex = parent;
    }
    return nodes.at(elemIndex).parent;
}

Define helper functions to improve readability

All of the parent and grandparent code gives the function above, even in its rewritten version, some clutter. We can greatly improve readability with the use of two helper functions:

S parentOf(S elemIndex) const {
    return nodes.at(elemIndex).parent;
}
// sets new parent and returns previous value
S setParent(S elemIndex, S newparent) const { 
    auto oldparent{parentOf(elemIndex)};
    nodes.at(elemIndex).parent = newparent;
    return oldparent;
}

Now we can rewrite getRepr like this:

S getRepr(S elemIndex) const {
    // Follow parent pointers until we reach a representative
    for (const auto curr = parentOf(elemIndex); curr != parentOf(curr); )  {
        // set parent = grandparent
        elemIndex = setParent(elemIndex, parentOf(curr));
    }
    return parentOf(elemIndex);
}

I find that easier to read and reason about and it's shorter.

Answer 2

Prefer using operator[] instead of at()

The member function at() performs bounds checking, and this comes at the cost of performance. Most of the time, you already know the indices you are using are valid. The exception is when someone can pass arbitrary values to public member functions such as areInSameSet(), so ideally you do the bounds checking at the top of such functions, if at all. Consider that the user of this class already knows the number of elements of the set, so it's unlikely they would ever call member functions with incorrect indices.

Use uint8_t for rank

In the comments you say that rank is never negative, but then you explicitly declare it to be signed. I would make it an uint8_t, as this tells you it's never going to be negative, and it's being used to hold an integer, not a character.

Answer 3

High-level design:

All the other answers went over the minutiae quite sufficiently, but this is where about 66% savings are easily achieved.

There are some questionable design-decisions, which complicate the code and might just bite you hard:

1. Any query of nodes will always try to compress the path.

   While it is a good idea not to redo the work needlessly, this means that the constant interface is neither reentrancy nor, as there is no internal locking, concurrency safe.

   Why not have a read-only interface and a mutable interface, letting the caller decide on his requirements? Also, add a function optimizing the set for querying.

2. rank is useless, as there is a better way to limit the depth of trees.

   • It bloats a Node by half.
   • The rank of an interior-node is irrelevant, as it is never exposed nor acted on.
   • The rank of a root node is a simple function of the count of its descendants: \$rank = \lceil ld(size)\rceil\$
   • Still, no need to calculate it, as only relative rank is needed for limiting depth, which is easily done directly comparing size.

3. Any Node always uses either size or parent. Combine them for saving half the space, if more nodes are needed let them choose a bigger type.

   Half the range for ids, half for sizes (without zero).

Answer 4

Why do you use S for size type? You can just use std::size_t

nondefault assignment operator with default copy constructor is unidiomatic, just let your assignment operator as DisjointSet& operator=(const DisjointSet& other) = default;

Setting getRepr as const but mutating parent inside looks very awkward, just make getRepr as non-const (and give a better function name)