122 lines
3.3 KiB
C++
122 lines
3.3 KiB
C++
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
#pragma once
|
|
|
|
#include <algorithm>
|
|
#include <cassert>
|
|
#include <queue>
|
|
#include <utility>
|
|
|
|
#include "common/compiler_util.h" // IWYU pragma: keep
|
|
|
|
namespace doris {
|
|
|
|
template <typename T, typename Sequence, typename Compare>
|
|
class SortingHeap {
|
|
public:
|
|
SortingHeap(const Compare& comp) : _comp(comp) {}
|
|
|
|
bool is_valid() const { return !_queue.empty(); }
|
|
|
|
T top() { return _queue.front(); }
|
|
|
|
size_t size() { return _queue.size(); }
|
|
|
|
bool empty() { return _queue.empty(); }
|
|
|
|
T& next_child() { return _queue[_next_child_index()]; }
|
|
|
|
void replace_top(T new_top) {
|
|
*_queue.begin() = new_top;
|
|
update_top();
|
|
}
|
|
|
|
void remove_top() {
|
|
std::pop_heap(_queue.begin(), _queue.end(), _comp);
|
|
_queue.pop_back();
|
|
_next_idx = 0;
|
|
}
|
|
|
|
void push(T cursor) {
|
|
_queue.emplace_back(cursor);
|
|
std::push_heap(_queue.begin(), _queue.end(), _comp);
|
|
_next_idx = 0;
|
|
}
|
|
|
|
Sequence&& sorted_seq() {
|
|
std::sort_heap(_queue.begin(), _queue.end(), _comp);
|
|
return std::move(_queue);
|
|
}
|
|
|
|
private:
|
|
Sequence _queue;
|
|
Compare _comp;
|
|
|
|
/// Cache comparison between first and second child if the order in queue has not been changed.
|
|
size_t _next_idx = 0;
|
|
|
|
size_t _next_child_index() {
|
|
if (_next_idx == 0) {
|
|
_next_idx = 1;
|
|
if (_queue.size() > 2 && _comp(_queue[1], _queue[2])) ++_next_idx;
|
|
}
|
|
|
|
return _next_idx;
|
|
}
|
|
|
|
void update_top() {
|
|
size_t size = _queue.size();
|
|
if (size < 2) return;
|
|
|
|
auto begin = _queue.begin();
|
|
|
|
size_t child_idx = _next_child_index();
|
|
auto child_it = begin + child_idx;
|
|
|
|
/// Check if we are in order.
|
|
if (_comp(*child_it, *begin)) return;
|
|
|
|
_next_idx = 0;
|
|
|
|
auto curr_it = begin;
|
|
auto top = *begin;
|
|
do {
|
|
/// We are not in heap-order, swap the parent with it's largest child.
|
|
*curr_it = *child_it;
|
|
curr_it = child_it;
|
|
|
|
// recompute the child based off of the updated parent
|
|
child_idx = 2 * child_idx + 1;
|
|
|
|
if (child_idx >= size) break;
|
|
|
|
child_it = begin + child_idx;
|
|
|
|
if ((child_idx + 1) < size && _comp(*child_it, *(child_it + 1))) {
|
|
/// Right child exists and is greater than left child.
|
|
++child_it;
|
|
++child_idx;
|
|
}
|
|
|
|
/// Check if we are in order.
|
|
} while (!(_comp(*child_it, top)));
|
|
*curr_it = top;
|
|
}
|
|
};
|
|
} // namespace doris
|