UnionFind.h

/**********************************************************************
 *
 * GEOS - Geometry Engine Open Source
 * http://geos.osgeo.org
 *
 * Copyright (C) 2020-2021 Daniel Baston
 *
 * This is free software; you can redistribute and/or modify it under
 * the terms of the GNU Lesser General Public Licence as published
 * by the Free Software Foundation.
 * See the COPYING file for more information.
 *
 **********************************************************************/
 
#pragma once
 
#include <algorithm>
#include <numeric>
#include <vector>
 
#include <geos/export.h>
#include <geos/operation/cluster/Clusters.h>
 
namespace geos {
namespace operation {
namespace cluster {
 
class GEOS_DLL UnionFind {
 
public:
    explicit UnionFind(size_t n) :
            clusters(n),
            sizes(n),
            num_clusters(n) {
        std::iota(clusters.begin(), clusters.end(), 0);
        std::fill(sizes.begin(), sizes.end(), 1);
    }
 
    // Are two elements in the same cluster?
    bool same(size_t i, size_t j) {
        return i == j || (find(i) == find(j));
    }
 
    // Are two elements in a different cluster?
    bool different(size_t i, size_t j) {
        return !same(i, j);
    }
 
    size_t find(size_t i) {
        size_t root = i;
 
        while(clusters[root] != root) {
            root = clusters[root];
        }
 
        while (i != root) {
            size_t next = clusters[i];
            clusters[i] = root;
            i = next;
        }
 
        return i;
    }
 
    void join(size_t i, size_t j) {
        auto a = find(i);
        auto b = find(j);
 
        if (a == b) {
            return;
        }
 
        if ((sizes[b] > sizes[a]) || (sizes[a] == sizes[b] && b <= a)) {
            std::swap(a, b);
        }
 
        clusters[a] = clusters[b];
        sizes[b] += sizes[a];
        sizes[a] = 0;
 
        num_clusters--;
    }
 
    size_t getNumClusters() const {
        return num_clusters;
    }
 
    template<typename T>
    void sortByCluster(T begin, T end) {
        std::sort(begin, end, [this](size_t a, size_t b) {
            return find(a) < find(b);
        });
    }
 
    Clusters getClusters();
 
    Clusters getClusters(std::vector<size_t> elems);
 
private:
    std::vector<size_t> clusters;
    std::vector<size_t> sizes;
    size_t num_clusters;
};
 
}
}
}