Skip to content

Commit 22c1be0

Browse files
authoredDec 1, 2021
Merge pull request #500 from gareth-nx/pr/quickselect
Selection algorithms
2 parents 3af3259 + 0f4e635 commit 22c1be0

9 files changed

+1122
-0
lines changed
 

‎doc/specs/stdlib_selection.md

Lines changed: 350 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,350 @@
1+
---
2+
title: Selection Procedures
3+
---
4+
5+
# The `stdlib_selection` module
6+
7+
[TOC]
8+
9+
## Overview of selection
10+
11+
Suppose you wish to find the value of the k-th smallest entry in an array of size N, or
12+
the index of that value. While it could be done by sorting the whole array
13+
using `[[stdlib_sorting(module):sort(interface)]]` or
14+
`[[stdlib_sorting(module):sort_index(interface)]]` from
15+
`[[stdlib_sorting(module)]]` and then finding the k-th entry, that would
16+
require O(N x LOG(N)) time. However selection of a single entry can be done in
17+
O(N) time, which is much faster for large arrays. This is useful, for example,
18+
to quickly find the median of an array, or some other percentile.
19+
20+
The Fortran Standard Library therefore provides a module, `stdlib_selection`,
21+
which implements selection algorithms.
22+
23+
## Overview of the module
24+
25+
The module `stdlib_selection` defines two generic subroutines:
26+
* `select` is used to find the k-th smallest entry of an array. The input
27+
array is also modified in-place, and on return will be partially sorted
28+
such that `all(array(1:k) <= array(k)))` and `all(array(k) <= array((k+1):size(array)))` is true.
29+
The user can optionally specify `left` and `right` indices to constrain the search
30+
for the k-th smallest value. This can be useful if you have previously called `select`
31+
to find a smaller or larger rank (that will have led to partial sorting of
32+
`array`, thus implying some constraints on the location).
33+
34+
* `arg_select` is used to find the index of the k-th smallest entry of an array.
35+
In this case the input array is not modified, but the user must provide an
36+
input index array with the same size as `array`, having indices that are a permutation of
37+
`1:size(array)`, which is modified instead. On return the index array is modified
38+
such that `all(array(index(1:k)) <= array(index(k)))` and `all(array(k) <= array(k+1:size(array)))`.
39+
The user can optionally specify `left` and `right` indices to constrain the search
40+
for the k-th smallest value. This can be useful if you have previously called `arg_select`
41+
to find a smaller or larger rank (that will have led to partial sorting of
42+
`index`, thus implying some constraints on the location).
43+
44+
45+
## `select` - find the k-th smallest value in an input array
46+
47+
### Status
48+
49+
Experimental
50+
51+
### Description
52+
53+
Returns the k-th smallest value of `array(:)`, and also partially sorts `array(:)`
54+
such that `all(array(1:k) <= array(k))` and `all(array(k) <= array((k+1):size(array)))`
55+
56+
### Syntax
57+
58+
`call [[stdlib_selection(module):select(interface)]]( array, k, kth_smallest [, left, right ] )`
59+
60+
### Class
61+
62+
Generic subroutine.
63+
64+
### Arguments
65+
66+
`array` : shall be a rank one array of any of the types:
67+
`integer(int8)`, `integer(int16)`, `integer(int32)`, `integer(int64)`,
68+
`real(sp)`, `real(dp)`, `real(xdp)`, `real(qp)`. It is an `intent(inout)` argument.
69+
70+
`k`: shall be a scalar with any of the types:
71+
`integer(int8)`, `integer(int16)`, `integer(int32)`, `integer(int64)`. It
72+
is an `intent(in)` argument. We search for the `k`-th smallest entry of `array(:)`.
73+
74+
`kth_smallest`: shall be a scalar with the same type as `array`. It is an
75+
`intent(out)` argument. On return it contains the k-th smallest entry of
76+
`array(:)`.
77+
78+
`left` (optional): shall be a scalar with the same type as `k`. It is an
79+
`intent(in)` argument. If specified then we assume the k-th smallest value is
80+
definitely contained in `array(left:size(array))`. If `left` is not present,
81+
the default is 1. This is typically useful if multiple calls to `select` are
82+
made, because the partial sorting of `array` implies constraints on where we
83+
need to search.
84+
85+
`right` (optional): shall be a scalar with the same type as `k`. It is an
86+
`intent(in)` argument. If specified then we assume the k-th smallest value is
87+
definitely contained in `array(1:right)`. If `right` is not present, the
88+
default is `size(array)`. This is typically useful if multiple calls to
89+
`select` are made, because the partial sorting of `array` implies constraints
90+
on where we need to search.
91+
92+
### Notes
93+
94+
Selection of a single value should have runtime of O(`size(array)`), so it is
95+
asymptotically faster than sorting `array` entirely. The test program at the
96+
end of this document shows that is the case.
97+
98+
The code does not support `NaN` elements in `array`; it will run, but there is
99+
no consistent interpretation given to the order of `NaN` entries of `array`
100+
compared to other entries.
101+
102+
`select` was derived from code in the Coretran library by Leon Foks,
103+
https://github.com/leonfoks/coretran. Leon Foks has given permission for the
104+
code here to be released under stdlib's MIT license.
105+
106+
### Example
107+
108+
```fortran
109+
program demo_select
110+
use stdlib_selection, only: select
111+
implicit none
112+
113+
real, allocatable :: array(:)
114+
real :: kth_smallest
115+
integer :: k, left, right
116+
117+
array = [3., 2., 7., 4., 5., 1., 4., -1.]
118+
119+
k = 2
120+
call select(array, k, kth_smallest)
121+
print*, kth_smallest ! print 1.0
122+
123+
k = 7
124+
! Due to the previous call to select, we know for sure this is in an
125+
! index >= 2
126+
call select(array, k, kth_smallest, left=2)
127+
print*, kth_smallest ! print 5.0
128+
129+
k = 6
130+
! Due to the previous two calls to select, we know for sure this is in
131+
! an index >= 2 and <= 7
132+
call select(array, k, kth_smallest, left=2, right=7)
133+
print*, kth_smallest ! print 4.0
134+
135+
end program demo_select
136+
```
137+
138+
## `arg_select` - find the index of the k-th smallest value in an input array
139+
140+
### Status
141+
142+
Experimental
143+
144+
### Description
145+
146+
Returns the index of the k-th smallest value of `array(:)`, and also partially sorts
147+
the index-array `indx(:)` such that `all(array(indx(1:k)) <= array(indx(k)))` and
148+
`all(array(indx(k)) <= array(indx((k+1):size(array))))`
149+
150+
### Syntax
151+
152+
`call [[stdlib_selection(module):arg_select(interface)]]( array, indx, k, kth_smallest [, left, right ] )`
153+
154+
### Class
155+
156+
Generic subroutine.
157+
158+
### Arguments
159+
160+
`array` : shall be a rank one array of any of the types:
161+
`integer(int8)`, `integer(int16)`, `integer(int32)`, `integer(int64)`,
162+
`real(sp)`, `real(dp)`, `real(xdp), `real(qp)`. It is an `intent(in)` argument. On input it is
163+
the array in which we search for the k-th smallest entry.
164+
165+
`indx`: shall be a rank one array with the same size as `array`, containing all integers
166+
from `1:size(array)` in any order. It is of any of the types:
167+
`integer(int8)`, `integer(int16)`, `integer(int32)`, `integer(int64)`. It is an
168+
`intent(inout)` argument. On return its elements will define a partial sorting of `array(:)` such that:
169+
`all( array(indx(1:k-1)) <= array(indx(k)) )` and `all(array(indx(k)) <= array(indx(k+1:size(array))))`.
170+
171+
`k`: shall be a scalar with the same type as `indx`. It is an `intent(in)`
172+
argument. We search for the `k`-th smallest entry of `array(:)`.
173+
174+
`kth_smallest`: a scalar with the same type as `indx`. It is an `intent(out)` argument,
175+
and on return it contains the index of the k-th smallest entry of `array(:)`.
176+
177+
`left` (optional): shall be a scalar with the same type as `k`. It is an `intent(in)`
178+
argument. If specified then we assume the k-th smallest value is definitely contained
179+
in `array(indx(left:size(array)))`. If `left` is not present, the default is 1.
180+
This is typically useful if multiple calls to `arg_select` are made, because
181+
the partial sorting of `indx` implies constraints on where we need to search.
182+
183+
`right` (optional): shall be a scalar with the same type as `k`. It is an `intent(in)`
184+
argument. If specified then we assume the k-th smallest value is definitely contained
185+
in `array(indx(1:right))`. If `right` is not present, the default is
186+
`size(array)`. This is typically useful if multiple calls to `arg_select` are
187+
made, because the reordering of `indx` implies constraints on where we need to
188+
search.
189+
190+
### Notes
191+
192+
`arg_select` does not modify `array`, unlike `select`.
193+
194+
The partial sorting of `indx` is not stable, i.e., indices that map to equal
195+
values of array may be reordered.
196+
197+
The code does not support `NaN` elements in `array`; it will run, but there is
198+
no consistent interpretation given to the order of `NaN` entries of `array`
199+
compared to other entries.
200+
201+
While it is essential that that `indx` contains a permutation of the integers `1:size(array)`,
202+
the code does not check for this. For example if `size(array) == 4`, then we could have
203+
`indx = [4, 2, 1, 3]` or `indx = [1, 2, 3, 4]`, but not `indx = [2, 1, 2, 4]`. It is the user's
204+
responsibility to avoid such errors.
205+
206+
Selection of a single value should have runtime of O(`size(array)`), so it is
207+
asymptotically faster than sorting `array` entirely. The test program at the end of
208+
these documents confirms that is the case.
209+
210+
`arg_select` was derived using code from the Coretran library by Leon Foks,
211+
https://github.com/leonfoks/coretran. Leon Foks has given permission for the
212+
code here to be released under stdlib's MIT license.
213+
214+
### Example
215+
216+
217+
```fortran
218+
program demo_arg_select
219+
use stdlib_selection, only: arg_select
220+
implicit none
221+
222+
real, allocatable :: array(:)
223+
integer, allocatable :: indx(:)
224+
integer :: kth_smallest
225+
integer :: k, left, right
226+
227+
array = [3., 2., 7., 4., 5., 1., 4., -1.]
228+
indx = [( k, k = 1, size(array) )]
229+
230+
k = 2
231+
call arg_select(array, indx, k, kth_smallest)
232+
print*, array(kth_smallest) ! print 1.0
233+
234+
k = 7
235+
! Due to the previous call to arg_select, we know for sure this is in an
236+
! index >= 2
237+
call arg_select(array, indx, k, kth_smallest, left=2)
238+
print*, array(kth_smallest) ! print 5.0
239+
240+
k = 6
241+
! Due to the previous two calls to arg_select, we know for sure this is in
242+
! an index >= 2 and <= 7
243+
call arg_select(array, indx, k, kth_smallest, left=2, right=7)
244+
print*, array(kth_smallest) ! print 4.0
245+
246+
end program demo_arg_select
247+
```
248+
249+
## Comparison with using `sort`
250+
251+
The following program compares the timings of `select` and `arg_select` for
252+
computing the median of an array, vs using `sort` from `stdlib`. In theory we
253+
should see a speed improvement with the selection routines which grows like
254+
LOG(size(`array`)).
255+
256+
```fortran
257+
program selection_vs_sort
258+
use stdlib_kinds, only: dp, sp, int64
259+
use stdlib_selection, only: select, arg_select
260+
use stdlib_sorting, only: sort
261+
implicit none
262+
263+
call compare_select_sort_for_median(1)
264+
call compare_select_sort_for_median(11)
265+
call compare_select_sort_for_median(101)
266+
call compare_select_sort_for_median(1001)
267+
call compare_select_sort_for_median(10001)
268+
call compare_select_sort_for_median(100001)
269+
270+
contains
271+
subroutine compare_select_sort_for_median(N)
272+
integer, intent(in) :: N
273+
274+
integer :: i, k, result_arg_select, indx(N), indx_local(N)
275+
real :: random_vals(N), local_random_vals(N)
276+
integer, parameter :: test_reps = 100
277+
integer(int64) :: t0, t1
278+
real :: result_sort, result_select
279+
integer(int64) :: time_sort, time_select, time_arg_select
280+
logical :: select_test_passed, arg_select_test_passed
281+
282+
! Ensure N is odd
283+
if(mod(N, 2) /= 1) stop
284+
285+
time_sort = 0
286+
time_select = 0
287+
time_arg_select = 0
288+
289+
select_test_passed = .true.
290+
arg_select_test_passed = .true.
291+
292+
indx = (/( i, i = 1, N) /)
293+
294+
k = (N+1)/2 ! Deliberate integer division
295+
296+
do i = 1, test_reps
297+
call random_number(random_vals)
298+
299+
! Compute the median with sorting
300+
local_random_vals = random_vals
301+
call system_clock(t0)
302+
call sort(local_random_vals)
303+
result_sort = local_random_vals(k)
304+
call system_clock(t1)
305+
time_sort = time_sort + (t1 - t0)
306+
307+
! Compute the median with selection, assuming N is odd
308+
local_random_vals = random_vals
309+
call system_clock(t0)
310+
call select(local_random_vals, k, result_select)
311+
call system_clock(t1)
312+
time_select = time_select + (t1 - t0)
313+
314+
! Compute the median with arg_select, assuming N is odd
315+
local_random_vals = random_vals
316+
indx_local = indx
317+
call system_clock(t0)
318+
call arg_select(local_random_vals, indx_local, k, result_arg_select)
319+
call system_clock(t1)
320+
time_arg_select = time_arg_select + (t1 - t0)
321+
322+
if(result_select /= result_sort) select_test_passed = .FALSE.
323+
if(local_random_vals(result_arg_select) /= result_sort) arg_select_test_passed = .FALSE.
324+
end do
325+
326+
print*, "select ; N=", N, '; ', merge('PASS', 'FAIL', select_test_passed), &
327+
'; Relative-speedup-vs-sort:', (1.0*time_sort)/(1.0*time_select)
328+
print*, "arg_select; N=", N, '; ', merge('PASS', 'FAIL', arg_select_test_passed), &
329+
'; Relative-speedup-vs-sort:', (1.0*time_sort)/(1.0*time_arg_select)
330+
331+
end subroutine
332+
333+
end program
334+
```
335+
336+
The results seem consistent with expectations when the `array` is large; the program prints:
337+
```
338+
select ; N= 1 ; PASS; Relative-speedup-vs-sort: 1.90928173
339+
arg_select; N= 1 ; PASS; Relative-speedup-vs-sort: 1.76875830
340+
select ; N= 11 ; PASS; Relative-speedup-vs-sort: 1.14835048
341+
arg_select; N= 11 ; PASS; Relative-speedup-vs-sort: 1.00794709
342+
select ; N= 101 ; PASS; Relative-speedup-vs-sort: 2.31012774
343+
arg_select; N= 101 ; PASS; Relative-speedup-vs-sort: 1.92877376
344+
select ; N= 1001 ; PASS; Relative-speedup-vs-sort: 4.24190664
345+
arg_select; N= 1001 ; PASS; Relative-speedup-vs-sort: 3.54580402
346+
select ; N= 10001 ; PASS; Relative-speedup-vs-sort: 5.61573362
347+
arg_select; N= 10001 ; PASS; Relative-speedup-vs-sort: 4.79348087
348+
select ; N= 100001 ; PASS; Relative-speedup-vs-sort: 7.28823519
349+
arg_select; N= 100001 ; PASS; Relative-speedup-vs-sort: 6.03007460
350+
```

‎src/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ set(fppFiles
1212
stdlib_linalg_diag.fypp
1313
stdlib_linalg_outer_product.fypp
1414
stdlib_optval.fypp
15+
stdlib_selection.fypp
1516
stdlib_sorting.fypp
1617
stdlib_sorting_ord_sort.fypp
1718
stdlib_sorting_sort.fypp

‎src/Makefile.manual

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ SRCFYPP = \
1313
stdlib_quadrature.fypp \
1414
stdlib_quadrature_trapz.fypp \
1515
stdlib_quadrature_simps.fypp \
16+
stdlib_selection.fypp \
1617
stdlib_random.fypp \
1718
stdlib_sorting.fypp \
1819
stdlib_sorting_ord_sort.fypp \
@@ -105,6 +106,8 @@ stdlib_quadrature_trapz.o: \
105106
stdlib_quadrature.o \
106107
stdlib_error.o \
107108
stdlib_kinds.o
109+
stdlib_selection.o: \
110+
stdlib_kinds.o
108111
stdlib_sorting.o: \
109112
stdlib_kinds.o \
110113
stdlib_string_type.o

‎src/stdlib_selection.fypp

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
#:include "common.fypp"
2+
! Specify kinds/types for the input array in select and arg_select
3+
#:set ARRAY_KINDS_TYPES = INT_KINDS_TYPES + REAL_KINDS_TYPES
4+
! The index arrays are of all INT_KINDS_TYPES
5+
6+
module stdlib_selection
7+
!
8+
! This code was modified from the "Coretran" implementation "quickSelect" by
9+
! Leon Foks, https://github.com/leonfoks/coretran/tree/master/src/sorting
10+
!
11+
! Leon Foks gave permission to release this code under stdlib's MIT license.
12+
! (https://github.com/fortran-lang/stdlib/pull/500#commitcomment-57418593)
13+
!
14+
15+
use stdlib_kinds
16+
17+
implicit none
18+
19+
private
20+
21+
public :: select, arg_select
22+
23+
interface select
24+
!! version: experimental
25+
!! ([Specification](..//page/specs/stdlib_selection.html#select-find-the-k-th-smallest-value-in-an-input-array))
26+
27+
#:for arraykind, arraytype in ARRAY_KINDS_TYPES
28+
#:for intkind, inttype in INT_KINDS_TYPES
29+
#:set name = rname("select", 1, arraytype, arraykind, intkind)
30+
module procedure ${name}$
31+
#:endfor
32+
#:endfor
33+
end interface
34+
35+
interface arg_select
36+
!! version: experimental
37+
!! ([Specification](..//page/specs/stdlib_selection.html#arg_select-find-the-index-of-the-k-th-smallest-value-in-an-input-array))
38+
#:for arraykind, arraytype in ARRAY_KINDS_TYPES
39+
#:for intkind, inttype in INT_KINDS_TYPES
40+
#:set name = rname("arg_select", 1, arraytype, arraykind, intkind)
41+
module procedure ${name}$
42+
#:endfor
43+
#:endfor
44+
end interface
45+
46+
contains
47+
48+
#:for arraykind, arraytype in ARRAY_KINDS_TYPES
49+
#:for intkind, inttype in INT_KINDS_TYPES
50+
#:set name = rname("select", 1, arraytype, arraykind, intkind)
51+
subroutine ${name}$(a, k, kth_smallest, left, right)
52+
!! select - select the k-th smallest entry in a(:).
53+
!!
54+
!! Partly derived from the "Coretran" implementation of
55+
!! quickSelect by Leon Foks, https://github.com/leonfoks/coretran
56+
!!
57+
${arraytype}$, intent(inout) :: a(:)
58+
!! Array in which we seek the k-th smallest entry.
59+
!! On output it will be partially sorted such that
60+
!! `all(a(1:(k-1)) <= a(k)) .and. all(a(k) <= a((k+1):size(a)))`
61+
!! is true.
62+
${inttype}$, intent(in) :: k
63+
!! We want the k-th smallest entry. E.G. `k=1` leads to
64+
!! `kth_smallest=min(a)`, and `k=size(a)` leads to
65+
!! `kth_smallest=max(a)`
66+
${arraytype}$, intent(out) :: kth_smallest
67+
!! On output contains the k-th smallest value of `a(:)`
68+
${inttype}$, intent(in), optional :: left, right
69+
!! If we know that:
70+
!! the k-th smallest entry of `a` is in `a(left:right)`
71+
!! and also that:
72+
!! `maxval(a(1:(left-1))) <= minval(a(left:right))`
73+
!! and:
74+
!! `maxval(a(left:right))) <= minval(a((right+1):size(a)))`
75+
!! then one or both bounds can be specified to narrow the search.
76+
!! The constraints are available if we have previously called the
77+
!! subroutine with different `k` (because of how `a(:)` becomes
78+
!! partially sorted, see documentation for `a(:)`).
79+
80+
${inttype}$ :: l, r, mid, iPivot
81+
integer, parameter :: ip = ${intkind}$
82+
83+
l = 1_ip
84+
if(present(left)) l = left
85+
r = size(a, kind=ip)
86+
if(present(right)) r = right
87+
88+
if(k < 1_ip .or. k > size(a, kind=ip) .or. l > r .or. l < 1_ip .or. &
89+
r > size(a, kind=ip)) then
90+
error stop "select must have 1 <= k <= size(a), and 1 <= left <= right <= size(a)";
91+
end if
92+
93+
searchk: do
94+
mid = l + ((r-l)/2_ip) ! Avoid (l+r)/2 which can cause overflow
95+
96+
call medianOf3(a, l, mid, r)
97+
call swap(a(l), a(mid))
98+
call partition(a, l, r, iPivot)
99+
100+
if (iPivot < k) then
101+
l = iPivot + 1_ip
102+
elseif (iPivot > k) then
103+
r = iPivot - 1_ip
104+
elseif (iPivot == k) then
105+
kth_smallest = a(k)
106+
return
107+
end if
108+
end do searchk
109+
110+
contains
111+
pure subroutine swap(a, b)
112+
${arraytype}$, intent(inout) :: a, b
113+
${arraytype}$ :: tmp
114+
tmp = a; a = b; b = tmp
115+
end subroutine
116+
117+
pure subroutine medianOf3(a, left, mid, right)
118+
${arraytype}$, intent(inout) :: a(:)
119+
${inttype}$, intent(in) :: left, mid, right
120+
if(a(right) < a(left)) call swap(a(right), a(left))
121+
if(a(mid) < a(left)) call swap(a(mid) , a(left))
122+
if(a(right) < a(mid) ) call swap(a(mid) , a(right))
123+
end subroutine
124+
125+
pure subroutine partition(array,left,right,iPivot)
126+
${arraytype}$, intent(inout) :: array(:)
127+
${inttype}$, intent(in) :: left, right
128+
${inttype}$, intent(out) :: iPivot
129+
${inttype}$ :: lo,hi
130+
${arraytype}$ :: pivot
131+
132+
pivot = array(left)
133+
lo = left
134+
hi=right
135+
do while (lo <= hi)
136+
do while (array(hi) > pivot)
137+
hi=hi-1_ip
138+
end do
139+
inner_lohi: do while (lo <= hi )
140+
if(array(lo) > pivot) exit inner_lohi
141+
lo=lo+1_ip
142+
end do inner_lohi
143+
if (lo <= hi) then
144+
call swap(array(lo),array(hi))
145+
lo=lo+1_ip
146+
hi=hi-1_ip
147+
end if
148+
end do
149+
call swap(array(left),array(hi))
150+
iPivot=hi
151+
end subroutine
152+
end subroutine
153+
#:endfor
154+
#:endfor
155+
156+
157+
#:for arraykind, arraytype in ARRAY_KINDS_TYPES
158+
#:for intkind, inttype in INT_KINDS_TYPES
159+
#:set name = rname("arg_select", 1, arraytype, arraykind, intkind)
160+
subroutine ${name}$(a, indx, k, kth_smallest, left, right)
161+
!! arg_select - find the index of the k-th smallest entry in `a(:)`
162+
!!
163+
!! Partly derived from the "Coretran" implementation of
164+
!! quickSelect by Leon Foks, https://github.com/leonfoks/coretran
165+
!!
166+
${arraytype}$, intent(in) :: a(:)
167+
!! Array in which we seek the k-th smallest entry.
168+
${inttype}$, intent(inout) :: indx(:)
169+
!! Array of indices into `a(:)`. Must contain each integer
170+
!! from `1:size(a)` exactly once. On output it will be partially
171+
!! sorted such that
172+
!! `all( a(indx(1:(k-1)))) <= a(indx(k)) ) .AND.
173+
!! all( a(indx(k)) <= a(indx( (k+1):size(a) )) )`.
174+
${inttype}$, intent(in) :: k
175+
!! We want index of the k-th smallest entry. E.G. `k=1` leads to
176+
!! `a(kth_smallest) = min(a)`, and `k=size(a)` leads to
177+
!! `a(kth_smallest) = max(a)`
178+
${inttype}$, intent(out) :: kth_smallest
179+
!! On output contains the index with the k-th smallest value of `a(:)`
180+
${inttype}$, intent(in), optional :: left, right
181+
!! If we know that:
182+
!! the k-th smallest entry of `a` is in `a(indx(left:right))`
183+
!! and also that:
184+
!! `maxval(a(indx(1:(left-1)))) <= minval(a(indx(left:right)))`
185+
!! and:
186+
!! `maxval(a(indx(left:right))) <= minval(a(indx((right+1):size(a))))`
187+
!! then one or both bounds can be specified to reduce the search
188+
!! time. These constraints are available if we have previously
189+
!! called the subroutine with a different `k` (due to the way that
190+
!! `indx(:)` becomes partially sorted, see documentation for `indx(:)`).
191+
192+
${inttype}$ :: l, r, mid, iPivot
193+
integer, parameter :: ip = ${intkind}$
194+
195+
l = 1_ip
196+
if(present(left)) l = left
197+
r = size(a, kind=ip)
198+
if(present(right)) r = right
199+
200+
if(size(a) /= size(indx)) then
201+
error stop "arg_select must have size(a) == size(indx)"
202+
end if
203+
204+
if(k < 1_ip .or. k > size(a, kind=ip) .or. l > r .or. l < 1_ip .or. &
205+
r > size(a, kind=ip)) then
206+
error stop "arg_select must have 1 <= k <= size(a), and 1 <= left <= right <= size(a)";
207+
end if
208+
209+
searchk: do
210+
mid = l + ((r-l)/2_ip) ! Avoid (l+r)/2 which can cause overflow
211+
212+
call arg_medianOf3(a, indx, l, mid, r)
213+
call swap(indx(l), indx(mid))
214+
call arg_partition(a, indx, l, r, iPivot)
215+
216+
if (iPivot < k) then
217+
l = iPivot + 1_ip
218+
elseif (iPivot > k) then
219+
r = iPivot - 1_ip
220+
elseif (iPivot == k) then
221+
kth_smallest = indx(k)
222+
return
223+
end if
224+
end do searchk
225+
226+
contains
227+
pure subroutine swap(a, b)
228+
${inttype}$, intent(inout) :: a, b
229+
${inttype}$ :: tmp
230+
tmp = a; a = b; b = tmp
231+
end subroutine
232+
233+
pure subroutine arg_medianOf3(a, indx, left, mid, right)
234+
${arraytype}$, intent(in) :: a(:)
235+
${inttype}$, intent(inout) :: indx(:)
236+
${inttype}$, intent(in) :: left, mid, right
237+
if(a(indx(right)) < a(indx(left))) call swap(indx(right), indx(left))
238+
if(a(indx(mid)) < a(indx(left))) call swap(indx(mid) , indx(left))
239+
if(a(indx(right)) < a(indx(mid)) ) call swap(indx(mid) , indx(right))
240+
end subroutine
241+
242+
pure subroutine arg_partition(array, indx, left,right,iPivot)
243+
${arraytype}$, intent(in) :: array(:)
244+
${inttype}$, intent(inout) :: indx(:)
245+
${inttype}$, intent(in) :: left, right
246+
${inttype}$, intent(out) :: iPivot
247+
${inttype}$ :: lo,hi
248+
${arraytype}$ :: pivot
249+
250+
pivot = array(indx(left))
251+
lo = left
252+
hi = right
253+
do while (lo <= hi)
254+
do while (array(indx(hi)) > pivot)
255+
hi=hi-1_ip
256+
end do
257+
inner_lohi: do while (lo <= hi )
258+
if(array(indx(lo)) > pivot) exit inner_lohi
259+
lo=lo+1_ip
260+
end do inner_lohi
261+
if (lo <= hi) then
262+
call swap(indx(lo),indx(hi))
263+
lo=lo+1_ip
264+
hi=hi-1_ip
265+
end if
266+
end do
267+
call swap(indx(left),indx(hi))
268+
iPivot=hi
269+
end subroutine
270+
end subroutine
271+
#:endfor
272+
#:endfor
273+
274+
end module
275+
276+

‎src/tests/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ add_subdirectory(io)
2121
add_subdirectory(linalg)
2222
add_subdirectory(logger)
2323
add_subdirectory(optval)
24+
add_subdirectory(selection)
2425
add_subdirectory(sorting)
2526
add_subdirectory(stats)
2627
add_subdirectory(string)

‎src/tests/Makefile.manual

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ all test clean::
1717
$(MAKE) -f Makefile.manual --directory=io $@
1818
$(MAKE) -f Makefile.manual --directory=logger $@
1919
$(MAKE) -f Makefile.manual --directory=optval $@
20+
$(MAKE) -f Makefile.manual --directory=selection $@
2021
$(MAKE) -f Makefile.manual --directory=sorting $@
2122
$(MAKE) -f Makefile.manual --directory=quadrature $@
2223
$(MAKE) -f Makefile.manual --directory=stats $@

‎src/tests/selection/CMakeLists.txt

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
### Pre-process: .fpp -> .f90 via Fypp
2+
3+
# Create a list of the files to be preprocessed
4+
set(
5+
fppFiles
6+
test_selection.fypp
7+
)
8+
9+
fypp_f90("${fyppFlags}" "${fppFiles}" outFiles)
10+
11+
ADDTEST(selection)

‎src/tests/selection/Makefile.manual

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
SRCFYPP = test_selection.fypp
2+
3+
SRCGEN = $(SRCFYPP:.fypp=.f90)
4+
5+
$(SRCGEN): %.f90: %.fypp ../../common.fypp
6+
fypp -I../.. $(FYPPFLAGS) $< $@
7+
8+
PROGS_SRC = $(SRCGEN)
9+
10+
include ../Makefile.manual.test.mk

‎src/tests/selection/test_selection.fypp

Lines changed: 469 additions & 0 deletions
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)
Please sign in to comment.