1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

fn median_of_three(list: &Vec<impl PartialOrd>, left: usize, right: usize) -> usize {
    let half_index = right / 2;

    let beginning = &list[left];
    let middle = &list[half_index];
    let end = &list[right];

    let comparison = |median, other_one, other_two| {
        median > other_one && median < other_two || median > other_two && median < other_one
    };

    if comparison(beginning, middle, end) {
        left
    } else if comparison(middle, end, beginning) {
        half_index
    } else if comparison(end, beginning, middle) {
        right
    } else {
        left
    }
}

// lomunto partition scheme
fn partition(
    list: &mut Vec<impl PartialOrd>,
    left: usize,
    right: usize,
    pivot_index: usize,
) -> usize {
    list.swap(pivot_index, right);

    let mut store_index = left;
    for i in left..right {
        if list[i] < list[right] {
            list.swap(i, store_index);
            store_index += 1
        }
    }
    list.swap(right, store_index);
    store_index
}

fn select<T: PartialOrd>(list: &mut Vec<T>, left: usize, right: usize, k: usize) -> &T {
    if left == right {
        &list[left]
    } else {
        let pivot_index = partition(list, left, right, median_of_three(&list, left, right));
        if k == pivot_index {
            &list[k]
        } else if k < pivot_index {
            select(list, left, pivot_index - 1, k)
        } else {
            select(list, pivot_index + 1, right, k)
        }
    }
}

/// Run quick select on vector to find kth smallest (0 based)
pub fn quick_select<T: PartialOrd>(array: &mut Vec<T>, k: usize) -> &T {
    select(array, 0, array.len() - 1, k)
}

#[cfg(test)]
mod tests {
    use crate::{partition, select, quick_select};

    #[test]
    fn partition_test() {
        let mut vector = vec![1, 1, -43, 0, -1, -1, -1];
        let len = vector.len() - 1;
        let result = partition(&mut vector, 0, len, 3);
        assert_eq!(vector, vec![-43, -1, -1, -1, 0, 1, 1]);
        assert_eq!(result, 4)
    }

    #[test]
    fn select_test() {
        let mut vector = vec![3, 2, 1, 6, 5, 4, 0];
        let len = vector.len() - 1;
        for i in 1..=len {
            let result = *select(&mut vector, 0, len, i);
            dbg!(&vector);
            assert_eq!(result, i)
        }
    }

    #[test]
    fn it_works() {
        let mut vector = vec![100, 2, 7, 1, 4, 43, 0];
        for i in 0..vector.len() {
            println!("{}", quick_select(&mut vector, i));
            dbg!(&vector);
        }

    }
}