aisd_lab/libselect/src/randomized_select.rs

use rand::{thread_rng, Rng};

use crate::{normal_select::NormalSelect, Select};

pub struct RandomizedSelect {
    comparisons: u64,
    swaps: u64,
    should_print: bool,
    normal: NormalSelect,
}

impl RandomizedSelect {

    fn rand_partition(&mut self, list: &mut Vec<u64>, lo: usize, hi: usize) -> usize {
        let pivot_index = thread_rng().gen_range(lo..hi);
        if self.should_print { println!("pivot_index: {pivot_index}"); }
        self.normal.partition(list, lo, hi, list[pivot_index])
    }

    fn _select(&mut self, list: &mut Vec<u64>, lo: usize, hi: usize, k: usize) -> (usize, u64) {
        if lo == hi {
            return (lo, list[lo]);
        }

        if self.should_print { println!("list: {:?}; lo: {lo}; hi: {hi}", list); }

        let r = self.rand_partition(list, lo, hi);
        if self.should_print { println!("r: {r}"); }

        self.comparisons += self.normal.num_comp();
        self.swaps += self.normal.num_swap();
        self.normal.reset_state();
        let i = r - lo + 1;
        if k == i {
            return (r, list[r]);
        } else if k < i {
            return self._select(list, lo, r - 1, k);
        } else {
            return self._select(list, r + 1, hi, k - i);
        }
    }
}

impl Select for RandomizedSelect {
    fn new(should_print: bool) -> Self {
        Self {
            comparisons: 0,
            swaps: 0,
            should_print,
            normal: NormalSelect::new(should_print),
        }
    }

    fn select_mut(&mut self, list: &mut Vec<u64>, order_statistic: usize) -> u64 {
        self._select(list, 0, list.len() - 1, order_statistic).1
    }

    fn num_comp(&self) -> u64 {
        self.comparisons
    }

    fn num_swap(&self) -> u64 {
        self.swaps
    }

    fn reset_state(&mut self) {
        self.comparisons = 0;
        self.swaps = 0;
        self.normal.reset_state();
    }
}