union-find.h revision 8fc5a7f51e62cb4ae44a27bdf4176d04adc80ede
1// union-find.h
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15//
16// \file Union-Find algorithm for dense sets of non-negative
17// integers. Implemented using disjoint tree forests with rank
18// heuristics and path compression.
19
20#ifndef __fst_union_find_inl_h__
21#define __fst_union_find_inl_h__
22
23#include <stack>
24#include <vector>
25
26namespace fst {
27
28// Union-Find algorithm for dense sets of non-negative integers
29// (exact type: T).
30template <class T>
31class UnionFind {
32 public:
33  // Ctor: creates a disjoint set forest for the range [0;max).
34  // 'fail' is a value indicating that an element hasn't been
35  // initialized using MakeSet(...). The upper bound of the range
36  // can be reset (increased) using MakeSet(...).
37  UnionFind(T max, T fail)
38      : parent_(max, fail), rank_(max), fail_(fail) { }
39
40  // Finds the representative of the set 'item' belongs to.
41  // Performs path compression if needed.
42  T FindSet(T item) {
43    if (item >= parent_.size()
44        || item == fail_
45        || parent_[item] == fail_) return fail_;
46
47    CHECK(exec_stack_.empty());
48    T *p = &parent_[item];
49    for (; *p != item; item = *p, p = &parent_[item]) {
50      exec_stack_.push(p);
51    }
52    for (; ! exec_stack_.empty(); exec_stack_.pop()) {
53      *exec_stack_.top() = *p;
54    }
55    return *p;
56  }
57
58  // Creates the (destructive) union of the sets x and y belong to.
59  void Union(T x, T y) {
60    Link(FindSet(x), FindSet(y));
61  }
62
63  // Initialization of an element: creates a singleton set containing
64  // 'item'. The range [0;max) is reset if item >= max.
65  T MakeSet(T item) {
66    if (item >= parent_.size()) {
67      // New value in parent_ should be initialized to fail_
68      parent_.resize(2 * item, fail_);
69      rank_.resize(2 * item);
70    }
71    parent_[item] = item;
72    return item;
73  }
74
75  // Initialization of all elements starting from 0 to max - 1 to distinct sets
76  void MakeAllSet(T max) {
77    parent_.resize(max);
78    for (T item = 0; item < max; ++item) {
79      parent_[item] = item;
80    }
81  }
82
83 private:
84  vector<T> parent_;      // Parent nodes.
85  vector<int> rank_;      // Rank of an element = min. depth in tree.
86  T fail_;                // Value indicating lookup failure.
87  stack<T*> exec_stack_;  // Used for path compression.
88
89  // Links trees rooted in 'x' and 'y'.
90  void Link(T x, T y) {
91    if (x == y) return;
92
93    if (rank_[x] > rank_[y]) {
94      parent_[y] = x;
95    } else {
96      parent_[x] = y;
97      if (rank_[x] == rank_[y]) {
98        ++rank_[y];
99      }
100    }
101  }
102  DISALLOW_EVIL_CONSTRUCTORS(UnionFind);
103};
104
105}
106
107#endif  // __fst_union_find_inl_h__
108