1/**************************************************************************/
2/* disjoint_set.h */
3/**************************************************************************/
4/* This file is part of: */
5/* GODOT ENGINE */
6/* https://godotengine.org */
7/**************************************************************************/
8/* Copyright (c) 2014-present Godot Engine contributors (see AUTHORS.md). */
9/* Copyright (c) 2007-2014 Juan Linietsky, Ariel Manzur. */
10/* */
11/* Permission is hereby granted, free of charge, to any person obtaining */
12/* a copy of this software and associated documentation files (the */
13/* "Software"), to deal in the Software without restriction, including */
14/* without limitation the rights to use, copy, modify, merge, publish, */
15/* distribute, sublicense, and/or sell copies of the Software, and to */
16/* permit persons to whom the Software is furnished to do so, subject to */
17/* the following conditions: */
18/* */
19/* The above copyright notice and this permission notice shall be */
20/* included in all copies or substantial portions of the Software. */
21/* */
22/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, */
23/* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF */
24/* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. */
25/* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY */
26/* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, */
27/* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE */
28/* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */
29/**************************************************************************/
30
31#ifndef DISJOINT_SET_H
32#define DISJOINT_SET_H
33
34#include "core/templates/rb_map.h"
35#include "core/templates/vector.h"
36
37/* This DisjointSet class uses Find with path compression and Union by rank */
38template <typename T, class H = HashMapHasherDefault, class C = HashMapComparatorDefault<T>, class AL = DefaultAllocator>
39class DisjointSet {
40 struct Element {
41 T object;
42 Element *parent = nullptr;
43 int rank = 0;
44 };
45
46 typedef HashMap<T, Element *, H, C> MapT;
47
48 MapT elements;
49
50 Element *get_parent(Element *element);
51
52 _FORCE_INLINE_ Element *insert_or_get(T object);
53
54public:
55 ~DisjointSet();
56
57 _FORCE_INLINE_ void insert(T object) { (void)insert_or_get(object); }
58
59 void create_union(T a, T b);
60
61 void get_representatives(Vector<T> &out_roots);
62
63 void get_members(Vector<T> &out_members, T representative);
64};
65
66/* FUNCTIONS */
67
68template <typename T, class H, class C, class AL>
69DisjointSet<T, H, C, AL>::~DisjointSet() {
70 for (KeyValue<T, Element *> &E : elements) {
71 memdelete_allocator<Element, AL>(E.value);
72 }
73}
74
75template <typename T, class H, class C, class AL>
76typename DisjointSet<T, H, C, AL>::Element *DisjointSet<T, H, C, AL>::get_parent(Element *element) {
77 if (element->parent != element) {
78 element->parent = get_parent(element->parent);
79 }
80
81 return element->parent;
82}
83
84template <typename T, class H, class C, class AL>
85typename DisjointSet<T, H, C, AL>::Element *DisjointSet<T, H, C, AL>::insert_or_get(T object) {
86 typename MapT::Iterator itr = elements.find(object);
87 if (itr != nullptr) {
88 return itr->value;
89 }
90
91 Element *new_element = memnew_allocator(Element, AL);
92 new_element->object = object;
93 new_element->parent = new_element;
94 elements.insert(object, new_element);
95
96 return new_element;
97}
98
99template <typename T, class H, class C, class AL>
100void DisjointSet<T, H, C, AL>::create_union(T a, T b) {
101 Element *x = insert_or_get(a);
102 Element *y = insert_or_get(b);
103
104 Element *x_root = get_parent(x);
105 Element *y_root = get_parent(y);
106
107 // Already in the same set
108 if (x_root == y_root) {
109 return;
110 }
111
112 // Not in the same set, merge
113 if (x_root->rank < y_root->rank) {
114 SWAP(x_root, y_root);
115 }
116
117 // Merge y_root into x_root
118 y_root->parent = x_root;
119 if (x_root->rank == y_root->rank) {
120 ++x_root->rank;
121 }
122}
123
124template <typename T, class H, class C, class AL>
125void DisjointSet<T, H, C, AL>::get_representatives(Vector<T> &out_representatives) {
126 for (KeyValue<T, Element *> &E : elements) {
127 Element *element = E.value;
128 if (element->parent == element) {
129 out_representatives.push_back(element->object);
130 }
131 }
132}
133
134template <typename T, class H, class C, class AL>
135void DisjointSet<T, H, C, AL>::get_members(Vector<T> &out_members, T representative) {
136 typename MapT::Iterator rep_itr = elements.find(representative);
137 ERR_FAIL_NULL(rep_itr);
138
139 Element *rep_element = rep_itr->value;
140 ERR_FAIL_COND(rep_element->parent != rep_element);
141
142 for (KeyValue<T, Element *> &E : elements) {
143 Element *parent = get_parent(E.value);
144 if (parent == rep_element) {
145 out_members.push_back(E.key);
146 }
147 }
148}
149
150#endif // DISJOINT_SET_H
151