Calculating the median and mode of a list of integers

Question

Going through the Rust book for the second time as I've been wrapped up in Typescript for a while and I want to build a personal project in Rust.

In chapter 8.3, there's an exercise:

Given a list of integers, use a vector and return the median (when sorted, the value in the middle position) and mode (the value that occurs most often; a hash map will be helpful here) of the list.

I wrote the following functions:

fn main() {
    let mut v = vec![10, -1, 9, -2];
    println!("{}", exercise_median_of_vector(&mut v));

    let v = vec![1, 2, 3, 1];
    println!("{}", exercise_mode_of_vector(&v));
}

fn exercise_median_of_vector(v: &mut Vec<i32>) -> i32 {
    v.sort();

    let half = v.len() / 2;

    if half & 1 == 1 {
        v[half + 1]
    } else {
        v[half]
    }
}

fn exercise_mode_of_vector(vector: &Vec<i32>) -> i32 {
    struct Highest {
        key: i32,
        count: i32,
    }

    let mut mode = HashMap::new();
    let mut highest = Highest { key: 0, count: 0 };

    for key in vector {
        let count = mode.entry(key).or_insert(0);
        *count += 1;

        if *count > highest.count {
            highest.count = *count;
            highest.key = *key;
        }
    }

    highest.key
}

And whilst moving on, I thought it might be a good idea to get some input on this.

One of the gripes I have with the code is that the exercise_median_of_vector function mutates the input vector. How would I go about doing this immutably and efficiently?

Another thought I had was, in the mode function, to put the check inside an impl block; I wondered whether this would unnecessarily increase the memory footprint and/or performance. Any thoughts?

Answer 1

Your Code Is Buggy. Test It!

Try this testcase:

    assert_eq!(exercise_median_of_vector(&mut Vec::from([1])), 1);

Always write test cases for your code!

In a case like [2,4], you return the value 4 as the median. But, by the generally-accepted definition, the median is 3.

This gets tricky when you’re talking about returning the median of [1,2] as an i32, since i32 can only return a whole number. But, here, I’d consider the median of a sorted vector v of even length n (greater than 0) to be (v[n/2 - 1] + v[n/2])/2. I think that’s more in the spirit of the exercise.

Speaking of which, you don’t test for empty vectors. Your options here are to use, well, an Option<i32> and return None when handed an empty vector, or to panic. (Or I guess you could return a f64 that can represent “and a half,” or a nonexistent median as NaN.) If you’re going to panic, I’d suggest it be with something like

assert!(n > 0); // Maybe even with a panic message.

or

if (n == 0) {
    unimplemented!();
}

which is much more helpful to a maintainer, and at least shows you’ve thought about it.

Do You Really Mean &mut?

Right now, the input parameter is passed as a &mut Vec<i32>. The function alters the borrowed vector by sorting it and then moves an element out of it. By the Principle of Least Surprise, we would not expect a function that finds the median of a vector to permute it. Also, this gets in the way of fluent interfaces and makes you write &mut in front of arguments.

If you are going to modify v, including by moving the result out of it, you should take ownership of it, and remove &mut from the type. In that case, you can do anything you want to it because the borrow checker will prevent the caller from using it after.

You should only borrow v if you are going to leave it in a usable state afterward. In that case, you would want to work on a clone of v that you own.

The Mode Algorithm

You chose to implement a hash table with Key and Value types of i32. This could fail for a vector with billions of zeroes. The value type (that is, the count) should be usize. The key type should indeed be the Item type of the inputs, so that looks good.

You don’t need the Highest type for what could just be two local variables. It’s a little surprising to me that, in this function to calculate a mode, mode is a hash map and the result is highest.key.

Make Sure Twiddling Bits Actually Helps

In the current implementation, you have an if half & 1 == 1 control structure, but it’s currently equivalent to, v[(v.len() + 1)/2] (which is not correct).

if you change this to use the more common definition of median for even-length lists, you really do need the if block. In that case, half & 1 and half % 2 will optimize to equivalent code.

(I’d also probably name half, middle, but there’s nothing actually wrong with it. It is half of something.)

Consider Generics

This is what I think you were getting at when you brought up impl, although I’m not sure I understand your specific question.

The original exercise said, “Given a list of integers,” which you interpreted as a Vec<i32>. That’s a reasonable interpretation, given that the last section but one you just finished reading in the Rust Book was titled, “Storing Lists of Values with Vectors.”

But there are other types of lists and other types of integers, and you might want to support them. Since you want to solve the problem on your own, but you haven’t gotten to the necessary chapter in the Rust book yet, here’s a skeleton:

use std::ops::Add;
use std::ops::Div;

pub fn generic_median<T : IntoIterator>(range: T) -> Option<T::Item>
where
    T::Item : Add<T::Item, Output = T::Item>, // Needed to add items
    T::Item : Copy, // Needed to move items.
    T::Item : Ord, // Needed to sort a vector of items
    T::Item : From<u8>, // Needed for from(2)
    T::Item : Div<T::Item, Output = T::Item> // Needed to divide items by T::Item::from(2)
{
    let mut v : Vec<T::Item> =
        range.into_iter()
             .collect();
    // This part has been left as an exercise for the reader.
}

Here are some testcases for it:

use std::collections::LinkedList;

assert_eq!(generic_median([3,2,1,2]), Some(2));
assert_eq!{generic_median(LinkedList::from([2u16,1,3])), Some(2)};
assert_eq!(generic_median(Vec::from([1i64, 0i64])), Some(0i64));
assert_eq!(generic_median(Vec::<i32>::new()), None);

Some things that you might have expected to work, but don’t:

• Floating-point numbers, because vec::sort() does not work on them. The exercise specified integers. (This is because they are not fully-ordered; they could be NaN.)
• The i8 type

Consider Traits

You asked about an impl block, although I’m not sure which check you meant. You would use those with traits, and there are several advantages to writing traits for these, such as:

trait Median {
    type Output;
    fn median(self: Self) -> Option<Self::Output>;
}

For example, now you can use fluent syntax (v.median()). More importantly, you can write:

impl<Item> Median for Vec<Item>
where
    Item : Add<Item, Output = Item>, // Needed to add items
    Item : Copy, // Needed to move items.
    Item : Div<Item, Output = Item>, // Needed to divide items
    Item : From<u8>, // Needed to divide by Item::from(2)
    Item : Ord // Needed for .sort()
{
    type Output = Item;
    
    fn median(self: Self) -> Option<Self::Output> {
        let mut v = self;
        // Does this remind you of anything?
    }
}

The previous implementation, called with a Vec, made an unnecessary copy of the Vec, when it could have sorted the input in place. This is because the code had to work on generic containers, and Rust does not currently let you override the generic median to optimize for one type.

Some other containers have fast methods to move their contents into vectors, such as (the admittedly-contrived example) String::into_bytes, which we would want to use if optimizing for them.

This also works better if we want to change the signature to

fn median(self: &self) -> Option<Output>

And get a .median() that does not modify its input. In many cases, we would want to write something like

let mut v : Vec<Item> = input.iter()
                             .cloned()
                             .collect();

But .iter() is not currently a trait in Rust; it is duck-typed. So we could not write a generic version of this code in Rust (as of 2023). But we can use it in trait implementations.

Answer 2

Following that beautiful criticism, I hobbled through and came up with the following code. I've not yet reached the traits or generics sections yet, so I'm as of yet unsure as to whether or not this is 100% the right way to use them. I've also not reached the testing section yet either.

Taking notes from the Median example @Davislor so helpfully provided, I wrote the Mode trait too.

I think that it would be nice for me to revisit this and write a more cohesive unit test suite later on down the line but hopefully these assertions tick the box for now :)

Once I've got a bit more know-how, I'll try to do the borrowing version of this too.

use std::{
    collections::HashMap,
    hash::Hash,
    ops::{Add, Div},
};

trait Median {
    type Output;
    fn median(self: Self) -> Option<Self::Output>;
}

impl<Item> Median for Vec<Item>
where
    Item: Add<Item, Output = Item>, // Needed to add items
    Item: Copy,                     // Needed to move items.
    Item: Div<i32, Output = Item>,  // Needed to divide i32s
    Item: Div<Item, Output = Item>, // Needed to divide items
    Item: From<u8>,                 // Needed for Item::from(2)
    Item: Ord,                      // Needed for .sort()
{
    type Output = Item;

    fn median(self: Self) -> Option<Self::Output> {
        let mut v = self;

        match v.len() {
            0 => None,
            1 => Some(v[0]),
            _ => {
                let middle = v.len() / 2;
                v.sort();

                if v.len() % 2 == 0 {
                    Some((v[middle + 1] + v[middle]) / 2)
                } else {
                    Some(v[middle])
                }
            }
        }
    }
}

trait Mode {
    type Output;
    fn mode(self: Self) -> Option<Self::Output>;
}

impl<Item> Mode for Vec<Item>
where
    Item: From<Item>,
    Item: Copy,
    Item: Eq,
    Item: Hash,
    Item: From<u8>,
{
    type Output = Item;

    fn mode(self: Self) -> Option<Self::Output> {
        if self.len() == 0 {
            return None;
        }

        let mut item_counter = HashMap::new();

        let mut highest_count_thus_far: usize = 0;
        let mut to_return = self[0];

        for key in self {
            let count = item_counter.entry(key).or_insert(0);
            *count += 1;

            if *count > highest_count_thus_far {
                highest_count_thus_far = *count;
                to_return = key;
            }
        }

        Some(to_return)
    }
}

fn main() {
    assert_eq!(Median::median(vec![1, 2, 3]), Some(2));
    assert_eq!(Median::median(vec![1]), Some(1));
    assert_eq!(Median::median(v), None);

    let v: Vec<i32> = vec![];
    assert_eq!(Mode::mode(v), None);
    assert_eq!(Mode::mode(vec![1]), Some(1));
    assert_eq!(Mode::mode(vec![3,1,2,1]), Some(1));

    let v: Vec<u8> = vec![3,1,2,1];
    assert_eq!(Mode::mode(v), Some(1));

}
```