1+ /* *
2+ * MIT License
3+ *
4+ * Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
5+ *
6+ * Permission is hereby granted, free of charge, to any person obtaining a copy
7+ * of this software and associated documentation files (the "Software"), to deal
8+ * in the Software without restriction, including without limitation the rights
9+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10+ * copies of the Software, and to permit persons to whom the Software is
11+ * furnished to do so, subject to the following conditions:
12+ *
13+ * The above copyright notice and this permission notice shall be included in all
14+ * copies or substantial portions of the Software.
15+ *
16+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22+ * SOFTWARE.
23+ * */
24+
25+ #ifndef UNIFIEDCACHE_TOP_N_HEAP_H
26+ #define UNIFIEDCACHE_TOP_N_HEAP_H
27+
28+ #include < algorithm>
29+ #include < cstdint>
30+
31+ namespace UC {
32+
33+ template <typename T, typename Compare = std::less<T>>
34+ class TopNHeap {
35+ public:
36+ using ValueType = T;
37+ using SizeType = uint32_t ;
38+ using ConstRef = const T&;
39+
40+ private:
41+ using IndexType = uint32_t ;
42+ std::vector<ValueType> val_{};
43+ std::vector<IndexType> idx_{};
44+ SizeType capacity_{0 };
45+ SizeType size_{0 };
46+ Compare cmp_{};
47+
48+ public:
49+ explicit TopNHeap (const SizeType capacity) noexcept (
50+ std::is_nothrow_default_constructible_v<Compare>)
51+ : capacity_{capacity} {
52+ val_.resize (capacity);
53+ idx_.resize (capacity);
54+ }
55+ TopNHeap (const TopNHeap&) = delete ;
56+ TopNHeap (const TopNHeap&&) = delete ;
57+ TopNHeap& operator =(const TopNHeap&) = delete ;
58+ TopNHeap& operator =(const TopNHeap&&) = delete ;
59+ ~TopNHeap () { Clear (); }
60+
61+ SizeType Size () const noexcept { return size_; }
62+ SizeType Capacity () const noexcept { return capacity_; }
63+ bool Empty () const noexcept { return size_ == 0 ; }
64+
65+ void Push (ConstRef value) noexcept {
66+ if (size_ < capacity_) {
67+ val_[size_] = value;
68+ idx_[size_] = size_;
69+ SiftUp (size_);
70+ size_++;
71+ return ;
72+ }
73+ if (cmp_ (val_[idx_.front ()], value)) {
74+ val_[idx_.front ()] = value;
75+ SiftDown (0 );
76+ }
77+ }
78+ ConstRef Top () const noexcept { return val_[idx_.front ()]; }
79+ void Pop () noexcept {
80+ idx_[0 ] = idx_[--size_];
81+ if (size_) { SiftDown (0 ); }
82+ }
83+ void Clear () noexcept { size_ = 0 ; }
84+ private:
85+ static IndexType Parent (IndexType i) noexcept { return (i - 1 ) / 2 ; }
86+ static IndexType Left (IndexType i) noexcept { return 2 * i + 1 ; }
87+ static IndexType Right (IndexType i) noexcept { return 2 * i + 2 ; }
88+ void SiftUp (IndexType i) noexcept {
89+ auto pos = i;
90+ while (pos > 0 ) {
91+ auto p = Parent (pos);
92+ if (!cmp_ (val_[idx_[pos]], val_[idx_[p]])) { break ; }
93+ std::swap (idx_[pos], idx_[p]);
94+ pos = p;
95+ }
96+ }
97+ void SiftDown (IndexType i) noexcept {
98+ auto pos = i;
99+ for (;;) {
100+ auto l = Left (pos);
101+ auto r = Right (pos);
102+ auto best = pos;
103+ if (l < size_ && cmp_ (val_[idx_[l]], val_[idx_[best]])) { best = l; }
104+ if (r < size_ && cmp_ (val_[idx_[r]], val_[idx_[best]])) { best = r; }
105+ if (best == pos) { break ; }
106+ std::swap (idx_[pos], idx_[best]);
107+ pos = best;
108+ }
109+ }
110+ };
111+
112+ template <typename T, size_t N, typename Compare = std::less<T>>
113+ class TopNFixedHeap : public TopNHeap <T, Compare> {
114+ public:
115+ TopNFixedHeap () : TopNHeap<T, Compare>{N} {}
116+ };
117+
118+ } // namespace UC
119+
120+ #endif
0 commit comments