main.rs

#![allow(non_snake_case)]
use rand::seq::SliceRandom;
use std::time::Instant;
use std::ops::Deref;

pub struct NonEmptyMaxHeap<T> {
    elements: Vec<T>,
}

impl<T: Ord + Copy> NonEmptyMaxHeap<T> {
    pub fn init(root: T) -> Self {
        NonEmptyMaxHeap { elements: vec![root] }
    }

    pub fn insert(&mut self, element: T) {
        let mut index = self.elements.len();
        self.elements.push(element);

        while index > 0 {
            let parentIndex = (index - 1) / 2;
            let parent = self.elements[parentIndex];

            if !(parent < element) { break }

            self.elements[index] = parent;
            index = parentIndex;
        }

        self.elements[index] = element;
    }

    pub fn root(&self) -> T {
        self.elements[0]
    }

    pub fn replaceRoot(&mut self, element: T) {
        let mut index = 0;
        let mut childIndex = 2 * index + 1;

        while childIndex < self.elements.len() {
            let mut child = self.elements[childIndex];
            let rightIndex = childIndex + 1;

            if rightIndex < self.elements.len() {
                let right = self.elements[rightIndex];

                if child < right {
                    child = right;
                    childIndex = rightIndex;
                }
            }

            if !(element < child) { break }

            self.elements[index] = child;
            index = childIndex;
            childIndex = 2 * index + 1;
        }

        self.elements[index] = element;
    }
}

trait Smallest<T> {
    fn smallest(&self, n: usize) -> Vec<T>;
}

impl<T, I> Smallest<T> for I where T: Ord + Copy, I: Deref<Target = [T]> {
    fn smallest(&self, n: usize) -> Vec<T> {
        let mut iterator = self.iter().cloned();
        let first = match iterator.next() {
            Some(first) => first,
            None => return Vec::new(),
        };

        let mut heap = NonEmptyMaxHeap::init(first);
        let mut heapSize = 1;

        while heapSize < n {
            if let Some(element) = iterator.next() {
                heap.insert(element);
                heapSize += 1;
            } else {
                break;
            }
        }

        while let Some(element) = iterator.next() {
            if element < heap.root() {
                heap.replaceRoot(element);
            }
        }

        heap.elements.sort_unstable();
        heap.elements
    }
}

fn main() {
    let n = 1_000_000;
    let k = 1000;

    let mut rng = rand::thread_rng();
    let mut vec: Vec<_> = (0..n).collect();
    vec.shuffle(&mut rng);

    let start = Instant::now();
    let smallest = vec.smallest(k);
    println!("{:?}", start.elapsed());
    assert!(smallest.into_iter().eq(0..k));
    assert!(vec.len() == n);
}