xref: /freebsd/contrib/llvm-project/libcxx/include/__algorithm/nth_element.h (revision 1db9f3b21e39176dd5b67cf8ac378633b172463e)
1 //===----------------------------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef _LIBCPP___ALGORITHM_NTH_ELEMENT_H
10 #define _LIBCPP___ALGORITHM_NTH_ELEMENT_H
11 
12 #include <__algorithm/comp.h>
13 #include <__algorithm/comp_ref_type.h>
14 #include <__algorithm/iterator_operations.h>
15 #include <__algorithm/sort.h>
16 #include <__assert>
17 #include <__config>
18 #include <__debug_utils/randomize_range.h>
19 #include <__iterator/iterator_traits.h>
20 #include <__utility/move.h>
21 
22 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
23 #  pragma GCC system_header
24 #endif
25 
26 _LIBCPP_BEGIN_NAMESPACE_STD
27 
28 template <class _Compare, class _RandomAccessIterator>
29 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 bool __nth_element_find_guard(
30     _RandomAccessIterator& __i, _RandomAccessIterator& __j, _RandomAccessIterator __m, _Compare __comp) {
31   // manually guard downward moving __j against __i
32   while (true) {
33     if (__i == --__j) {
34       return false;
35     }
36     if (__comp(*__j, *__m)) {
37       return true; // found guard for downward moving __j, now use unguarded partition
38     }
39   }
40 }
41 
42 template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
43 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
44 // NOLINTNEXTLINE(readability-function-cognitive-complexity)
45 __nth_element(
46     _RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp) {
47   using _Ops = _IterOps<_AlgPolicy>;
48 
49   // _Compare is known to be a reference type
50   typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
51   const difference_type __limit = 7;
52   while (true) {
53     if (__nth == __last)
54       return;
55     difference_type __len = __last - __first;
56     switch (__len) {
57     case 0:
58     case 1:
59       return;
60     case 2:
61       if (__comp(*--__last, *__first))
62         _Ops::iter_swap(__first, __last);
63       return;
64     case 3: {
65       _RandomAccessIterator __m = __first;
66       std::__sort3<_AlgPolicy, _Compare>(__first, ++__m, --__last, __comp);
67       return;
68     }
69     }
70     if (__len <= __limit) {
71       std::__selection_sort<_AlgPolicy, _Compare>(__first, __last, __comp);
72       return;
73     }
74     // __len > __limit >= 3
75     _RandomAccessIterator __m   = __first + __len / 2;
76     _RandomAccessIterator __lm1 = __last;
77     unsigned __n_swaps          = std::__sort3<_AlgPolicy, _Compare>(__first, __m, --__lm1, __comp);
78     // *__m is median
79     // partition [__first, __m) < *__m and *__m <= [__m, __last)
80     // (this inhibits tossing elements equivalent to __m around unnecessarily)
81     _RandomAccessIterator __i = __first;
82     _RandomAccessIterator __j = __lm1;
83     // j points beyond range to be tested, *__lm1 is known to be <= *__m
84     // The search going up is known to be guarded but the search coming down isn't.
85     // Prime the downward search with a guard.
86     if (!__comp(*__i, *__m)) // if *__first == *__m
87     {
88       // *__first == *__m, *__first doesn't go in first part
89       if (std::__nth_element_find_guard<_Compare>(__i, __j, __m, __comp)) {
90         _Ops::iter_swap(__i, __j);
91         ++__n_swaps;
92       } else {
93         // *__first == *__m, *__m <= all other elements
94         // Partition instead into [__first, __i) == *__first and *__first < [__i, __last)
95         ++__i; // __first + 1
96         __j = __last;
97         if (!__comp(*__first, *--__j)) { // we need a guard if *__first == *(__last-1)
98           while (true) {
99             if (__i == __j) {
100               return; // [__first, __last) all equivalent elements
101             } else if (__comp(*__first, *__i)) {
102               _Ops::iter_swap(__i, __j);
103               ++__n_swaps;
104               ++__i;
105               break;
106             }
107             ++__i;
108           }
109         }
110         // [__first, __i) == *__first and *__first < [__j, __last) and __j == __last - 1
111         if (__i == __j) {
112           return;
113         }
114         while (true) {
115           while (!__comp(*__first, *__i)) {
116             ++__i;
117             _LIBCPP_ASSERT_UNCATEGORIZED(
118                 __i != __last,
119                 "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
120           }
121           do {
122             _LIBCPP_ASSERT_UNCATEGORIZED(
123                 __j != __first,
124                 "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
125             --__j;
126           } while (__comp(*__first, *__j));
127           if (__i >= __j)
128             break;
129           _Ops::iter_swap(__i, __j);
130           ++__n_swaps;
131           ++__i;
132         }
133         // [__first, __i) == *__first and *__first < [__i, __last)
134         // The first part is sorted,
135         if (__nth < __i) {
136           return;
137         }
138         // __nth_element the second part
139         // std::__nth_element<_Compare>(__i, __nth, __last, __comp);
140         __first = __i;
141         continue;
142       }
143     }
144     ++__i;
145     // j points beyond range to be tested, *__lm1 is known to be <= *__m
146     // if not yet partitioned...
147     if (__i < __j) {
148       // known that *(__i - 1) < *__m
149       while (true) {
150         // __m still guards upward moving __i
151         while (__comp(*__i, *__m)) {
152           ++__i;
153           _LIBCPP_ASSERT_UNCATEGORIZED(
154               __i != __last,
155               "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
156         }
157         // It is now known that a guard exists for downward moving __j
158         do {
159           _LIBCPP_ASSERT_UNCATEGORIZED(
160               __j != __first,
161               "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
162           --__j;
163         } while (!__comp(*__j, *__m));
164         if (__i >= __j)
165           break;
166         _Ops::iter_swap(__i, __j);
167         ++__n_swaps;
168         // It is known that __m != __j
169         // If __m just moved, follow it
170         if (__m == __i)
171           __m = __j;
172         ++__i;
173       }
174     }
175     // [__first, __i) < *__m and *__m <= [__i, __last)
176     if (__i != __m && __comp(*__m, *__i)) {
177       _Ops::iter_swap(__i, __m);
178       ++__n_swaps;
179     }
180     // [__first, __i) < *__i and *__i <= [__i+1, __last)
181     if (__nth == __i)
182       return;
183     if (__n_swaps == 0) {
184       // We were given a perfectly partitioned sequence.  Coincidence?
185       if (__nth < __i) {
186         // Check for [__first, __i) already sorted
187         __j = __m = __first;
188         while (true) {
189           if (++__j == __i) {
190             // [__first, __i) sorted
191             return;
192           }
193           if (__comp(*__j, *__m)) {
194             // not yet sorted, so sort
195             break;
196           }
197           __m = __j;
198         }
199       } else {
200         // Check for [__i, __last) already sorted
201         __j = __m = __i;
202         while (true) {
203           if (++__j == __last) {
204             // [__i, __last) sorted
205             return;
206           }
207           if (__comp(*__j, *__m)) {
208             // not yet sorted, so sort
209             break;
210           }
211           __m = __j;
212         }
213       }
214     }
215     // __nth_element on range containing __nth
216     if (__nth < __i) {
217       // std::__nth_element<_Compare>(__first, __nth, __i, __comp);
218       __last = __i;
219     } else {
220       // std::__nth_element<_Compare>(__i+1, __nth, __last, __comp);
221       __first = ++__i;
222     }
223   }
224 }
225 
226 template <class _AlgPolicy, class _RandomAccessIterator, class _Compare>
227 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __nth_element_impl(
228     _RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare& __comp) {
229   if (__nth == __last)
230     return;
231 
232   std::__debug_randomize_range<_AlgPolicy>(__first, __last);
233 
234   std::__nth_element<_AlgPolicy, __comp_ref_type<_Compare> >(__first, __nth, __last, __comp);
235 
236   std::__debug_randomize_range<_AlgPolicy>(__first, __nth);
237   if (__nth != __last) {
238     std::__debug_randomize_range<_AlgPolicy>(++__nth, __last);
239   }
240 }
241 
242 template <class _RandomAccessIterator, class _Compare>
243 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
244 nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp) {
245   std::__nth_element_impl<_ClassicAlgPolicy>(std::move(__first), std::move(__nth), std::move(__last), __comp);
246 }
247 
248 template <class _RandomAccessIterator>
249 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
250 nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last) {
251   std::nth_element(std::move(__first), std::move(__nth), std::move(__last), __less<>());
252 }
253 
254 _LIBCPP_END_NAMESPACE_STD
255 
256 #endif // _LIBCPP___ALGORITHM_NTH_ELEMENT_H
257