中文
devexamples

Top K Values with BinaryHeap

Keep the largest K values from a stream in a bounded min-heap using Reverse, then return them sorted from highest to lowest.

RustData Integrity

Code

snippet
use std::cmp::Reverse;
use std::collections::BinaryHeap;

fn top_k(values: &[i64], k: usize) -> Vec<i64> {
    if k == 0 {
        return Vec::new();
    }
    let mut heap: BinaryHeap<Reverse<i64>> = BinaryHeap::with_capacity(k + 1);
    for &value in values {
        heap.push(Reverse(value));
        if heap.len() > k {
            heap.pop();
        }
    }
    let mut out: Vec<i64> = heap.into_iter().map(|Reverse(value)| value).collect();
    out.sort_unstable();
    out.reverse();
    out
}

fn main() {
    let values = [42, 7, 19, 3, 88, 23, 56, 91, 5];
    println!("{:?}", top_k(&values, 3));
}

Related Examples