93 lines
2.4 KiB
Rust
93 lines
2.4 KiB
Rust
use std::cmp::Ordering;
|
|
|
|
fn partition_by<T, F: FnMut(&T, &T) -> Ordering>
|
|
(slice: &mut [T], start: isize, end: isize, compare: &mut F) -> isize {
|
|
|
|
let mut i = start;
|
|
for j in start + 1 .. end + 1 {
|
|
if compare(&slice[j as usize], &slice[start as usize]) != Ordering::Greater {
|
|
i += 1;
|
|
slice.swap(i as usize, j as usize);
|
|
}
|
|
}
|
|
slice.swap(start as usize, i as usize);
|
|
i
|
|
|
|
}
|
|
|
|
fn quick_sort_by<T, F: FnMut(&T, &T) -> Ordering>
|
|
(slice: &mut [T], start: isize, end: isize, compare: &mut F) {
|
|
|
|
if start < end {
|
|
let index = partition_by(slice, start, end, compare);
|
|
quick_sort_by(slice, start, index - 1, compare);
|
|
quick_sort_by(slice, index + 1, end, compare);
|
|
}
|
|
|
|
}
|
|
|
|
pub fn partial_sort_by<T, F: FnMut(&T, &T) -> Ordering>
|
|
(slice: &mut [T], k: usize, compare: &mut F) {
|
|
|
|
let mut start = 0;
|
|
let mut end = slice.len() as isize - 1;
|
|
|
|
while end > start {
|
|
|
|
let index = partition_by(slice, start, end, compare);
|
|
let rank = index + 1;
|
|
|
|
if rank >= k as isize {
|
|
end = index - 1;
|
|
} else if index - start > end - index {
|
|
quick_sort_by(slice, index + 1, end, compare);
|
|
end = index - 1;
|
|
} else {
|
|
quick_sort_by(slice, start, index - 1, compare);
|
|
start = index + 1;
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
/// Sorts the k first elements of the slice.
|
|
///
|
|
/// ```
|
|
/// # extern crate partial_sort;
|
|
/// # use partial_sort::partial_sort;
|
|
/// let mut v = vec![5, 7, 4, 2, 8, 6, 1, 9, 0, 3];
|
|
/// partial_sort(&mut v[..], 3);
|
|
/// assert_eq!(v[0], 0);
|
|
/// assert_eq!(v[1], 1);
|
|
/// assert_eq!(v[2], 2);
|
|
/// ```
|
|
///
|
|
/// The implementation is based on [this gist](https://gist.github.com/mesuvash/4095565).
|
|
pub fn partial_sort<T: Ord>(slice: &mut [T], k: usize) {
|
|
partial_sort_by(slice, k, &mut |x, y| x.cmp(y));
|
|
}
|
|
|
|
pub fn partial_sort_by_key<T, O: Ord, F: FnMut(&T) -> O>(slice: &mut [T], k: usize, key: &mut F) {
|
|
partial_sort_by(slice, k, &mut |x, y| key(x).cmp(&key(y)));
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
|
|
use partial_sort;
|
|
|
|
#[test]
|
|
fn partial_sort_test_1() {
|
|
let mut v = vec![5, 7, 4, 2, 8, 6, 1, 9, 0, 3];
|
|
partial_sort(&mut v[..], 7);
|
|
assert_eq!(v[0], 0);
|
|
assert_eq!(v[1], 1);
|
|
assert_eq!(v[2], 2);
|
|
assert_eq!(v[3], 3);
|
|
assert_eq!(v[4], 4);
|
|
assert_eq!(v[5], 5);
|
|
assert_eq!(v[6], 6);
|
|
}
|
|
|
|
}
|