xref: /llvm-project/libcxx/include/__algorithm/nth_element.h (revision 64184b4af02095745047fc620a81646664395bd2)
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 <__config>
13 #include <__algorithm/comp.h>
14 #include <__algorithm/comp_ref_type.h>
15 #include <__algorithm/sort.h>
16 #include <__iterator/iterator_traits.h>
17 #include <__utility/swap.h>
18 
19 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
20 #pragma GCC system_header
21 #endif
22 
23 _LIBCPP_BEGIN_NAMESPACE_STD
24 
25 template<class _Compare, class _RandomAccessIterator>
26 _LIBCPP_CONSTEXPR_AFTER_CXX11 bool
27 __nth_element_find_guard(_RandomAccessIterator& __i, _RandomAccessIterator& __j,
28                          _RandomAccessIterator __m, _Compare __comp)
29 {
30     // manually guard downward moving __j against __i
31     while (true) {
32         if (__i == --__j) {
33             return false;
34         }
35         if (__comp(*__j, *__m)) {
36             return true;  // found guard for downward moving __j, now use unguarded partition
37         }
38     }
39 }
40 
41 template <class _Compare, class _RandomAccessIterator>
42 _LIBCPP_CONSTEXPR_AFTER_CXX11 void
43 __nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp)
44 {
45     // _Compare is known to be a reference type
46     typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
47     const difference_type __limit = 7;
48     while (true)
49     {
50         if (__nth == __last)
51             return;
52         difference_type __len = __last - __first;
53         switch (__len)
54         {
55         case 0:
56         case 1:
57             return;
58         case 2:
59             if (__comp(*--__last, *__first))
60                 swap(*__first, *__last);
61             return;
62         case 3:
63             {
64             _RandomAccessIterator __m = __first;
65             _VSTD::__sort3<_Compare>(__first, ++__m, --__last, __comp);
66             return;
67             }
68         }
69         if (__len <= __limit)
70         {
71             _VSTD::__selection_sort<_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 = _VSTD::__sort3<_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 (_VSTD::__nth_element_find_guard<_Compare>(__i, __j, __m, __comp)) {
90                 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                             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                     while (__comp(*__first, *--__j))
118                         ;
119                     if (__i >= __j)
120                         break;
121                     swap(*__i, *__j);
122                     ++__n_swaps;
123                     ++__i;
124                 }
125                 // [__first, __i) == *__first and *__first < [__i, __last)
126                 // The first part is sorted,
127                 if (__nth < __i) {
128                     return;
129                 }
130                 // __nth_element the second part
131                 // _VSTD::__nth_element<_Compare>(__i, __nth, __last, __comp);
132                 __first = __i;
133                 continue;
134             }
135         }
136         ++__i;
137         // j points beyond range to be tested, *__lm1 is known to be <= *__m
138         // if not yet partitioned...
139         if (__i < __j)
140         {
141             // known that *(__i - 1) < *__m
142             while (true)
143             {
144                 // __m still guards upward moving __i
145                 while (__comp(*__i, *__m))
146                     ++__i;
147                 // It is now known that a guard exists for downward moving __j
148                 while (!__comp(*--__j, *__m))
149                     ;
150                 if (__i >= __j)
151                     break;
152                 swap(*__i, *__j);
153                 ++__n_swaps;
154                 // It is known that __m != __j
155                 // If __m just moved, follow it
156                 if (__m == __i)
157                     __m = __j;
158                 ++__i;
159             }
160         }
161         // [__first, __i) < *__m and *__m <= [__i, __last)
162         if (__i != __m && __comp(*__m, *__i))
163         {
164             swap(*__i, *__m);
165             ++__n_swaps;
166         }
167         // [__first, __i) < *__i and *__i <= [__i+1, __last)
168         if (__nth == __i)
169             return;
170         if (__n_swaps == 0)
171         {
172             // We were given a perfectly partitioned sequence.  Coincidence?
173             if (__nth < __i)
174             {
175                 // Check for [__first, __i) already sorted
176                 __j = __m = __first;
177                 while (true) {
178                     if (++__j == __i) {
179                         // [__first, __i) sorted
180                         return;
181                     }
182                     if (__comp(*__j, *__m)) {
183                         // not yet sorted, so sort
184                         break;
185                     }
186                     __m = __j;
187                 }
188             }
189             else
190             {
191                 // Check for [__i, __last) already sorted
192                 __j = __m = __i;
193                 while (true) {
194                     if (++__j == __last) {
195                         // [__i, __last) sorted
196                         return;
197                     }
198                     if (__comp(*__j, *__m)) {
199                         // not yet sorted, so sort
200                         break;
201                     }
202                     __m = __j;
203                 }
204             }
205         }
206         // __nth_element on range containing __nth
207         if (__nth < __i)
208         {
209             // _VSTD::__nth_element<_Compare>(__first, __nth, __i, __comp);
210             __last = __i;
211         }
212         else
213         {
214             // _VSTD::__nth_element<_Compare>(__i+1, __nth, __last, __comp);
215             __first = ++__i;
216         }
217     }
218 }
219 
220 template <class _RandomAccessIterator, class _Compare>
221 inline _LIBCPP_INLINE_VISIBILITY _LIBCPP_CONSTEXPR_AFTER_CXX17
222 void
223 nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp)
224 {
225     typedef typename __comp_ref_type<_Compare>::type _Comp_ref;
226     _VSTD::__nth_element<_Comp_ref>(__first, __nth, __last, __comp);
227 }
228 
229 template <class _RandomAccessIterator>
230 inline _LIBCPP_INLINE_VISIBILITY _LIBCPP_CONSTEXPR_AFTER_CXX17
231 void
232 nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last)
233 {
234     _VSTD::nth_element(__first, __nth, __last, __less<typename iterator_traits<_RandomAccessIterator>::value_type>());
235 }
236 
237 _LIBCPP_END_NAMESPACE_STD
238 
239 #endif // _LIBCPP___ALGORITHM_NTH_ELEMENT_H
240