1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
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 // This simple class finds the top n elements of an incrementally provided set
17 // of elements which you push one at a time. If the number of elements exceeds
18 // n, the lowest elements are incrementally dropped. At the end you get
19 // a vector of the top elements sorted in descending order (through Extract() or
20 // ExtractNondestructive()), or a vector of the top elements but not sorted
21 // (through ExtractUnsorted() or ExtractUnsortedNondestructive()).
22 //
23 // The value n is specified in the constructor. If there are p elements pushed
24 // altogether:
25 // The total storage requirements are O(min(n, p)) elements
26 // The running time is O(p * log(min(n, p))) comparisons
27 // If n is a constant, the total storage required is a constant and the running
28 // time is linear in p.
29 //
30 // NOTE(zhifengc): There is a way to do this in O(min(n, p)) storage and O(p)
31 // runtime. The basic idea is to repeatedly fill up a buffer of 2 * n elements,
32 // discarding the lowest n elements whenever the buffer is full using a linear-
33 // time median algorithm. This may have better performance when the input
34 // sequence is partially sorted.
35 //
36 // NOTE(zhifengc): This class should be redesigned to avoid reallocating a
37 // vector for each Extract.
38
39 #ifndef TENSORFLOW_LIB_GTL_TOP_N_H_
40 #define TENSORFLOW_LIB_GTL_TOP_N_H_
41
42 #include <stddef.h>
43 #include <algorithm>
44 #include <functional>
45 #include <string>
46 #include <vector>
47
48 #include "tensorflow/core/platform/logging.h"
49
50 namespace tensorflow {
51 namespace gtl {
52
53 // Cmp is an stl binary predicate. Note that Cmp is the "greater" predicate,
54 // not the more commonly used "less" predicate.
55 //
56 // If you use a "less" predicate here, the TopN will pick out the bottom N
57 // elements out of the ones passed to it, and it will return them sorted in
58 // ascending order.
59 //
60 // TopN is rule-of-zero copyable and movable if its members are.
61 template <class T, class Cmp = std::greater<T> >
62 class TopN {
63 public:
64 // The TopN is in one of the three states:
65 //
66 // o UNORDERED: this is the state an instance is originally in,
67 // where the elements are completely orderless.
68 //
69 // o BOTTOM_KNOWN: in this state, we keep the invariant that there
70 // is at least one element in it, and the lowest element is at
71 // position 0. The elements in other positions remain
72 // unsorted. This state is reached if the state was originally
73 // UNORDERED and a peek_bottom() function call is invoked.
74 //
75 // o HEAP_SORTED: in this state, the array is kept as a heap and
76 // there are exactly (limit_+1) elements in the array. This
77 // state is reached when at least (limit_+1) elements are
78 // pushed in.
79 //
80 // The state transition graph is at follows:
81 //
82 // peek_bottom() (limit_+1) elements
83 // UNORDERED --------------> BOTTOM_KNOWN --------------------> HEAP_SORTED
84 // | ^
85 // | (limit_+1) elements |
86 // +-----------------------------------------------------------+
87
88 enum State { UNORDERED, BOTTOM_KNOWN, HEAP_SORTED };
89 using UnsortedIterator = typename std::vector<T>::const_iterator;
90
91 // 'limit' is the maximum number of top results to return.
TopN(size_t limit)92 explicit TopN(size_t limit) : TopN(limit, Cmp()) {}
TopN(size_t limit,const Cmp & cmp)93 TopN(size_t limit, const Cmp &cmp) : limit_(limit), cmp_(cmp) {}
94
limit()95 size_t limit() const { return limit_; }
96
97 // Number of elements currently held by this TopN object. This
98 // will be no greater than 'limit' passed to the constructor.
size()99 size_t size() const { return std::min(elements_.size(), limit_); }
100
empty()101 bool empty() const { return size() == 0; }
102
103 // If you know how many elements you will push at the time you create the
104 // TopN object, you can call reserve to preallocate the memory that TopN
105 // will need to process all 'n' pushes. Calling this method is optional.
reserve(size_t n)106 void reserve(size_t n) { elements_.reserve(std::min(n, limit_ + 1)); }
107
108 // Push 'v'. If the maximum number of elements was exceeded, drop the
109 // lowest element and return it in 'dropped' (if given). If the maximum is not
110 // exceeded, 'dropped' will remain unchanged. 'dropped' may be omitted or
111 // nullptr, in which case it is not filled in.
112 // Requires: T is CopyAssignable, Swappable
push(const T & v)113 void push(const T &v) { push(v, nullptr); }
push(const T & v,T * dropped)114 void push(const T &v, T *dropped) { PushInternal(v, dropped); }
115
116 // Move overloads of push.
117 // Requires: T is MoveAssignable, Swappable
push(T && v)118 void push(T &&v) { // NOLINT(build/c++11)
119 push(std::move(v), nullptr);
120 }
push(T && v,T * dropped)121 void push(T &&v, T *dropped) { // NOLINT(build/c++11)
122 PushInternal(std::move(v), dropped);
123 }
124
125 // Peeks the bottom result without calling Extract()
126 const T &peek_bottom();
127
128 // Extract the elements as a vector sorted in descending order. The caller
129 // assumes ownership of the vector and must delete it when done. This is a
130 // destructive operation. The only method that can be called immediately
131 // after Extract() is Reset().
132 std::vector<T> *Extract();
133
134 // Similar to Extract(), but makes no guarantees the elements are in sorted
135 // order. As with Extract(), the caller assumes ownership of the vector and
136 // must delete it when done. This is a destructive operation. The only
137 // method that can be called immediately after ExtractUnsorted() is Reset().
138 std::vector<T> *ExtractUnsorted();
139
140 // A non-destructive version of Extract(). Copy the elements in a new vector
141 // sorted in descending order and return it. The caller assumes ownership of
142 // the new vector and must delete it when done. After calling
143 // ExtractNondestructive(), the caller can continue to push() new elements.
144 std::vector<T> *ExtractNondestructive() const;
145
146 // A non-destructive version of Extract(). Copy the elements to a given
147 // vector sorted in descending order. After calling
148 // ExtractNondestructive(), the caller can continue to push() new elements.
149 // Note:
150 // 1. The given argument must to be allocated.
151 // 2. Any data contained in the vector prior to the call will be deleted
152 // from it. After the call the vector will contain only the elements
153 // from the data structure.
154 void ExtractNondestructive(std::vector<T> *output) const;
155
156 // A non-destructive version of ExtractUnsorted(). Copy the elements in a new
157 // vector and return it, with no guarantees the elements are in sorted order.
158 // The caller assumes ownership of the new vector and must delete it when
159 // done. After calling ExtractUnsortedNondestructive(), the caller can
160 // continue to push() new elements.
161 std::vector<T> *ExtractUnsortedNondestructive() const;
162
163 // A non-destructive version of ExtractUnsorted(). Copy the elements into
164 // a given vector, with no guarantees the elements are in sorted order.
165 // After calling ExtractUnsortedNondestructive(), the caller can continue
166 // to push() new elements.
167 // Note:
168 // 1. The given argument must to be allocated.
169 // 2. Any data contained in the vector prior to the call will be deleted
170 // from it. After the call the vector will contain only the elements
171 // from the data structure.
172 void ExtractUnsortedNondestructive(std::vector<T> *output) const;
173
174 // Return an iterator to the beginning (end) of the container,
175 // with no guarantees about the order of iteration. These iterators are
176 // invalidated by mutation of the data structure.
unsorted_begin()177 UnsortedIterator unsorted_begin() const { return elements_.begin(); }
unsorted_end()178 UnsortedIterator unsorted_end() const { return elements_.begin() + size(); }
179
180 // Accessor for comparator template argument.
comparator()181 Cmp *comparator() { return &cmp_; }
182
183 // This removes all elements. If Extract() or ExtractUnsorted() have been
184 // called, this will put it back in an empty but useable state.
185 void Reset();
186
187 private:
188 template <typename U>
189 void PushInternal(U &&v, T *dropped); // NOLINT(build/c++11)
190
191 // elements_ can be in one of two states:
192 // elements_.size() <= limit_: elements_ is an unsorted vector of elements
193 // pushed so far.
194 // elements_.size() > limit_: The last element of elements_ is unused;
195 // the other elements of elements_ are an stl heap whose size is exactly
196 // limit_. In this case elements_.size() is exactly one greater than
197 // limit_, but don't use "elements_.size() == limit_ + 1" to check for
198 // that because you'll get a false positive if limit_ == size_t(-1).
199 std::vector<T> elements_;
200 size_t limit_; // Maximum number of elements to find
201 Cmp cmp_; // Greater-than comparison function
202 State state_ = UNORDERED;
203 };
204
205 // ----------------------------------------------------------------------
206 // Implementations of non-inline functions
207
208 template <class T, class Cmp>
209 template <typename U>
PushInternal(U && v,T * dropped)210 void TopN<T, Cmp>::PushInternal(U &&v, T *dropped) { // NOLINT(build/c++11)
211 if (limit_ == 0) {
212 if (dropped) *dropped = std::forward<U>(v); // NOLINT(build/c++11)
213 return;
214 }
215 if (state_ != HEAP_SORTED) {
216 elements_.push_back(std::forward<U>(v)); // NOLINT(build/c++11)
217 if (state_ == UNORDERED || cmp_(elements_.back(), elements_.front())) {
218 // Easy case: we just pushed the new element back
219 } else {
220 // To maintain the BOTTOM_KNOWN state, we need to make sure that
221 // the element at position 0 is always the smallest. So we put
222 // the new element at position 0 and push the original bottom
223 // element in the back.
224 // Warning: this code is subtle.
225 using std::swap;
226 swap(elements_.front(), elements_.back());
227 }
228 if (elements_.size() == limit_ + 1) {
229 // Transition from unsorted vector to a heap.
230 std::make_heap(elements_.begin(), elements_.end(), cmp_);
231 if (dropped) *dropped = std::move(elements_.front());
232 std::pop_heap(elements_.begin(), elements_.end(), cmp_);
233 state_ = HEAP_SORTED;
234 }
235 } else {
236 // Only insert the new element if it is greater than the least element.
237 if (cmp_(v, elements_.front())) {
238 // Store new element in the last slot of elements_. Remember from the
239 // comments on elements_ that this last slot is unused, so we don't
240 // overwrite anything useful.
241 elements_.back() = std::forward<U>(v); // NOLINT(build/c++11)
242
243 // stp::pop_heap() swaps elements_.front() and elements_.back() and
244 // rearranges elements from [elements_.begin(), elements_.end() - 1) such
245 // that they are a heap according to cmp_. Net effect: remove
246 // elements_.front() from the heap, and add the new element instead. For
247 // more info, see https://en.cppreference.com/w/cpp/algorithm/pop_heap.
248 std::pop_heap(elements_.begin(), elements_.end(), cmp_);
249 if (dropped) *dropped = std::move(elements_.back());
250 } else {
251 if (dropped) *dropped = std::forward<U>(v); // NOLINT(build/c++11)
252 }
253 }
254 }
255
256 template <class T, class Cmp>
peek_bottom()257 const T &TopN<T, Cmp>::peek_bottom() {
258 CHECK(!empty());
259 if (state_ == UNORDERED) {
260 // We need to do a linear scan to find out the bottom element
261 int min_candidate = 0;
262 for (size_t i = 1; i < elements_.size(); ++i) {
263 if (cmp_(elements_[min_candidate], elements_[i])) {
264 min_candidate = i;
265 }
266 }
267 // By swapping the element at position 0 and the minimal
268 // element, we transition to the BOTTOM_KNOWN state
269 if (min_candidate != 0) {
270 using std::swap;
271 swap(elements_[0], elements_[min_candidate]);
272 }
273 state_ = BOTTOM_KNOWN;
274 }
275 return elements_.front();
276 }
277
278 template <class T, class Cmp>
Extract()279 std::vector<T> *TopN<T, Cmp>::Extract() {
280 auto out = new std::vector<T>;
281 out->swap(elements_);
282 if (state_ != HEAP_SORTED) {
283 std::sort(out->begin(), out->end(), cmp_);
284 } else {
285 out->pop_back();
286 std::sort_heap(out->begin(), out->end(), cmp_);
287 }
288 return out;
289 }
290
291 template <class T, class Cmp>
ExtractUnsorted()292 std::vector<T> *TopN<T, Cmp>::ExtractUnsorted() {
293 auto out = new std::vector<T>;
294 out->swap(elements_);
295 if (state_ == HEAP_SORTED) {
296 // Remove the limit_+1'th element.
297 out->pop_back();
298 }
299 return out;
300 }
301
302 template <class T, class Cmp>
ExtractNondestructive()303 std::vector<T> *TopN<T, Cmp>::ExtractNondestructive() const {
304 auto out = new std::vector<T>;
305 ExtractNondestructive(out);
306 return out;
307 }
308
309 template <class T, class Cmp>
ExtractNondestructive(std::vector<T> * output)310 void TopN<T, Cmp>::ExtractNondestructive(std::vector<T> *output) const {
311 CHECK(output);
312 *output = elements_;
313 if (state_ != HEAP_SORTED) {
314 std::sort(output->begin(), output->end(), cmp_);
315 } else {
316 output->pop_back();
317 std::sort_heap(output->begin(), output->end(), cmp_);
318 }
319 }
320
321 template <class T, class Cmp>
ExtractUnsortedNondestructive()322 std::vector<T> *TopN<T, Cmp>::ExtractUnsortedNondestructive() const {
323 auto elements = new std::vector<T>;
324 ExtractUnsortedNondestructive(elements);
325 return elements;
326 }
327
328 template <class T, class Cmp>
ExtractUnsortedNondestructive(std::vector<T> * output)329 void TopN<T, Cmp>::ExtractUnsortedNondestructive(std::vector<T> *output) const {
330 CHECK(output);
331 *output = elements_;
332 if (state_ == HEAP_SORTED) {
333 // Remove the limit_+1'th element.
334 output->pop_back();
335 }
336 }
337
338 template <class T, class Cmp>
Reset()339 void TopN<T, Cmp>::Reset() {
340 elements_.clear();
341 state_ = UNORDERED;
342 }
343
344 } // namespace gtl
345 } // namespace tensorflow
346
347 #endif // TENSORFLOW_LIB_GTL_TOP_N_H_
348