Rust Unique Return Counts / Unique with Frequency

January 6, 2023

What is the fastest way to to get the unique elements in a vector and their count? Similar to numpy.unique(return_counts=True). The below becomes exceedingly slow as the array grows into the millions.

use std::collections::HashMap;
use itertools::Itertools;

fn main () {
    let kmers: Vec<u8> = vec![64, 64, 64, 65, 65, 65];
    let nodes: HashMap<u8, usize> = kmers
        .iter()
        .unique()
        .map(|kmer| {
            let count = kmers.iter().filter(|x| x == &kmer).count();
            (kmer.to_owned(), count)
        })
        .collect();
    println!("{:?}", nodes)   
}

>Solution :

You can use the entry API for this. The linked docs have a similar example to what you need, here it is modified to fit your case:

use std::collections::HashMap;

fn main () {
    let kmers: Vec<u8> = vec![64, 64, 64, 65, 65, 65];
    let mut nodes: HashMap<u8, usize> = HashMap::new();
    for n in kmers.iter() {
        nodes.entry(*n).and_modify(|count| *count += 1).or_insert(1);
    }

    println!("{:?}", nodes)   
}

playground

If you want the output to be sorted, you can use a BTreeMap instead.