2
\$\begingroup\$

Problem statement: A couple beginning a family decides to keep having children until they have at least one of either sex. Estimate the average number of children they will have via simulation. Also estimate the most common outcome. Assume that the probability \$p\$ of having a boy or girl is \$\frac{1}{2}\$.

This is one of my self-imposed challenges in Rust to become better at it. The problem was taken from Sedgewick Web Exercise 1.3.13

Here is my code:

use clap::Parser;
use rand::rngs::ThreadRng;
use rand::Rng;
use std::collections::HashMap;
use std::ops::RangeFrom;
use std::process::exit;

const VALID_TRIAL_NUMBERS: RangeFrom<u32> = 1..;

#[derive(Debug, Parser)]
struct Arguments {
    #[arg(index = 1)]
    number_of_trials: u32,
}

fn main() {
    let arguments = Arguments::parse();
    let number_of_trials: u32 = arguments.number_of_trials;
    let mut rng = rand::thread_rng();

    let Ok(frequencies) = simulate_many_families(number_of_trials, &mut rng) else {
        eprintln!("Number of trials must be at least {}.", VALID_TRIAL_NUMBERS.start);
        exit(1);
    };

    let average_children_number = calculate_average_children_number(number_of_trials, &frequencies);

    println!("Result of the simulation after {number_of_trials} trials:");
    println!("Average number of children is {average_children_number}.");

    if let Some(most_common_outcome) = find_most_common_outcome(&frequencies) {
        println!("The most common outcome is having {most_common_outcome} children.");
    } else {
        eprintln!("The frequencies HashMap is empty.");
    }
}

fn find_most_common_outcome(frequencies: &HashMap<u32, u32>) -> Option<u32> {
    let most_common_outcome: Option<u32> = frequencies
        .iter()
        .max_by_key(|(&_key, &value)| value)
        .map(|(&key, &_value)| key);
    most_common_outcome
}

fn calculate_average_children_number(
    number_of_trials: u32,
    frequencies: &HashMap<u32, u32>,
) -> f64 {
    let average_children_number: f64 = f64::round(
        f64::from(frequencies.iter().map(|(x, y)| x * y).sum::<u32>())
            / f64::from(number_of_trials),
    );
    average_children_number
}

fn simulate_many_families(
    number_of_trials: u32,
    rng: &mut ThreadRng,
) -> Result<HashMap<u32, u32>, String> {
    if !VALID_TRIAL_NUMBERS.contains(&number_of_trials) {
        return Err(format!(
            "Number of trials must be at least {}.",
            VALID_TRIAL_NUMBERS.start
        ));
    }

    let mut frequencies: HashMap<u32, u32> = HashMap::new();

    for _ in 0..number_of_trials {
        let number_of_children = simulate_one_family(rng);
        let count = frequencies.entry(number_of_children).or_insert(0);
        *count += 1;
    }

    Ok(frequencies)
}

fn simulate_one_family(rng: &mut ThreadRng) -> u32 {
    let mut number_of_boys: u32 = 0;
    let mut number_of_girls: u32 = 0;

    while number_of_boys < 1 || number_of_girls < 1 {
        let child: u32 = rng.gen_range(0..=1);

        match child {
            0 => number_of_boys += 1,
            1 => number_of_girls += 1,
            _ => (),
        }
    }

    number_of_boys + number_of_girls
}

Is there any way that I can improve my code?

\$\endgroup\$

1 Answer 1

3
\$\begingroup\$

This looks pretty good. It’s already been factored into functions that each do one thing.

As before, I’m going to suggest a way to do this with iterators and higher-order functions, because I’m weird and love functional programming. But first, let’s look at a way you might refactor the program.

Currently, you have a function

fn simulate_one_family(rng: &mut ThreadRng) -> u32

(Minor nit: I usually name a function like this according to the value it returns, and would name it simulate_one_family if it were primarily a side-effect function, but that’s a matter of style. It’s clear enough what this does.)

This is a great building block to start with, but pretend for the moment the team leader tells you, “We want to sum up the results of billions of trials, and a u32 could overflow, so change the return type to u64. We already have an iterator over infinite random bool values. Use that instead of ThreadRng. I’ll send you the documentation later.” that isn’t a big change to the interface:

fn simulate_one_family(rng: &mut impl Iterator<Item = bool>) -> u64

Inside the function, you currently have a while loop with two pieces of mutable state, number_of_boys and number_of_girls. You guessed it: I’m going to recommend you turn that into an iterator expression. This is a good candidate for Iterator::scan.

fn simulate_one_family(rng: &mut impl Iterator<Item = bool>) -> u64 {
    rng.scan(
        (0u64, 0u64),
        move |(number_of_girls, number_of_boys), is_boy| {
            if *number_of_girls > 0 && *number_of_boys > 0 {
                None
            } else if u64::MAX - *number_of_girls < *number_of_boys {
                panic!("Either the most-improbable thing ever has happened, or the program has a logic error.");
            } else if is_boy {
                *number_of_boys += 1;
                Some(*number_of_girls + *number_of_boys)
            } else {
                *number_of_girls +=1;
                Some(*number_of_girls + *number_of_boys)
            }
        },
    )
    .last()
    .expect("Logic error: There should have been at least two kids.")
}

This is a pretty close translation of your algorithm, with a few embellishments. Examining the second argument of scan first, it’s a closure |(number_of_girls, number_of_boys), is_boy| that returns an Option.

You can see where number_of_girls and number_of_boys went: they make up the state of the algorithm, and a scan takes the state as a single object. So they got packed into a tuple. This is actually passed by &mut instead of being moved.

A bigger change is that the closure in scan now gets a single random bit, is_boy, from the iterator rng. The code to generate that bit has been moved elsewhere, into the impl Iterator of rng’s type, whatever that is.

Back up to the first argument of scan, this is the initial value of the state: both number_of_girls and number_of_boys are initialized to zero.

The body of the closure is a conditional expression. A scan generates a sequence of values until it finishes with None, so the first test is whether to stop. This is equivalent to the loop condition of the original while. Next, I added a pro forma check that the families didn’t get so large, they overflow a u64, but you should be more worried that your computer will be hit by a meteor that was struck by lightning. The last two cases both increment the proper variable and add the total number of kids to the sequence. Finally, we take the last value of this sequence that isn’t None, and unwrap it, because we know the sequence cannot be empty. (There must be at least two kids before we stop.)

It has some advantages in terms of flexibility (you could, for example, feed a depterministic RNG, or some arbitrary non-random sequence of bits to it, without changing the code), without a real increase in complexity, although that might fall under YAGNI.

The main function is pretty different; The biggest change was taking out all the other helpers and ran multiple simulations using repeat_with, followed by take to limit it to a finite size. Your for loop will still work, but generally, collect on an iterator, will be more efficient than appending values one at a time to a mut Vec. Here, I just fed the iterator to .sum() instead of storing the values in memory at all.

I’d generally recommend using static single assignments and irrefutable patterns where possible. You’ll notice several necessary exceptions to this rule in the code. Some of the helpers in the original have been refactored into two or three lines in main.

Of course, I lied when I said we already had the iterator. It’s not in the rand crate you were using. We actually need to write it ourselves. That’s not the main focus here; a scan just needs its inputs in the form of an iterator. This consists of a struct to hold the state, a .next() function to get random bits, and From and Default to initialize it.

Finally, the problem specified simulating the results, but I’ll briefly mention the closed form solution. Let’s do a little combinatorics.

As I noted above, there have to be at least two children for there to be both a boy and a girl. There will be no more kids than that if this happens with the second child, plus one more if it has not happened until after the second child, plus one more than that if it has not happened until after the third, plus one more than that if it has not happened after the fourth, and so on.

Out of the 2**k possible outcomes of having k kids, only two don’t include a boy and a girl: k girls, or k boys. So, the probability of it not having happened after two kids is 2/2² = 1/2, the probability of it not having happened after three kids is 2/2³ = 1/4, and after k kids. is 1/2**(k - 1). The expected value is therefore 2 + 1/2 + 1/4 + 1/8 ... = 3. I use this to print how close to the theoretical expected value the results were.

Update: Branchless Code

The version above has an if is_boy test, where is_boy is an unpredictable branch. Those are hugely slow on some CPUs. Others can mitigate the misprediction penalty by using conditional moves.

More efficient is to write branchless code. Here is a version that removes the final branch, and calculates the number of girls and boys to add to the family from the value of the flag:

pub fn simulate_one_family(rng: &mut impl Iterator<Item = bool>) -> u64 {
    rng.scan(
        (0u64, 0u64),
        move |(number_of_girls, number_of_boys), is_boy| {
            if *number_of_girls > 0 && *number_of_boys > 0 {
                None
            } else if u64::MAX - *number_of_girls < *number_of_boys {
                panic!("Either the most-improbable thing ever has happened, or the program has a logic error.")
            } else {
                *number_of_girls += 1 - is_boy as u64;
                *number_of_boys += is_boy as u64;
                Some (*number_of_girls + *number_of_boys)
            }
        },
    )
    .last()
    .expect("Logic error: There should have been at least two kids.")
}

This generates better code, using sete and setne followed by add instead of cmove, on rustc 1.70.0 with -C opt-level=3 -C target-cpu=x86-64-v3. The relevant part of the loop, minus the RNG code, now looks like:

 test    rbx, rbx
 je      .LBB13_89
.LBB13_88:
 test    r14, r14
 jne     .LBB13_92
.LBB13_89:
 mov     r8, rbx
 add     r8, r14
 jb      .LBB13_90
 xor     r8d, r8d
 xor     r9d, r9d
 test    rdx, r10
 sete    r8b
 setne   r9b
 add     rbx, r8
 add     r14, r9
 lea     rdi, [r14, +, rbx]

In this assembly, rbx holds *number_of_girls, r14 holds *number_of_boys, test rdx, r10 checks is_boy and rdi holds the return value.

The first two branches check *number_of_girls > 0 and *number_of_boys > 0. The third is the overflow check. These should all be predictable by the CPU. The remaining code, to update *number_of_girls, *number_of_boys and the return value, has no branch instructions. This is surrounded by inlined iterator code.

Putting it All Together

extern crate rand;

/* Represents an iterator over infinite random `bool` values generatrd by
 * a (pseudo-)RNG.  A little more complicated than it needed to be, because
 * it uses every bit of entropy that the RNG returns.
 *
 * This is intended to be used through the traits it implements.  You should
 * not access the members directly.
 */
struct RandBoolIter<RNG: rand::RngCore> {
    rng: RNG,
    rand_bits: u64,
    bit_counter: u8,
}

/* Creates an iterator from a RNG.  The returned object will load 64 random
 * bits into its state the first time next() is called on it.
 */
impl<RNG: rand::RngCore> From<RNG> for RandBoolIter<RNG> {
    fn from(rng: RNG) -> RandBoolIter<RNG> {
        RandBoolIter {
            rng: rng,
            rand_bits: 0,
            bit_counter: u64::BITS as u8,
        }
    }
} 

// Like from, but uses the default RNG of its type.
impl<RNG: rand::RngCore + Default> RandBoolIter<RNG> {
    fn new() -> RandBoolIter<RNG> {
        Self::from(RNG::default())
    }
}

// Also allow a default RandBoolIter.
impl<RNG: rand::RngCore + Default> Default for RandBoolIter<RNG> {
    fn default() -> Self {
        Self::from(RNG::default())  
    }
}

impl<RNG: rand::RngCore> Iterator for RandBoolIter<RNG> {
    type Item = bool;

    fn next(&mut self) -> Option<bool> {
        if self.bit_counter == u64::BITS as u8 {
            self.rand_bits = self.rng.next_u64();
            self.bit_counter = 0;
        } else {
            self.bit_counter += 1;
        }

        Some(self.rand_bits & (1u64 << self.bit_counter) != 0)
    }
}

pub fn simulate_one_family(rng: &mut impl Iterator<Item = bool>) -> u64 {
    rng.scan(
        (0u64, 0u64),
        move |(number_of_girls, number_of_boys), is_boy| {
            if *number_of_girls > 0 && *number_of_boys > 0 {
                None
            } else if u64::MAX - *number_of_girls < *number_of_boys {
                panic!("Either the most-improbable thing ever has happened, or the program has a logic error.")
            } else {
                *number_of_girls += 1 - is_boy as u64;
                *number_of_boys += is_boy as u64;
                Some (*number_of_girls + *number_of_boys)
            }
        },
    )
    .last()
    .expect("Logic error: There should have been at least two kids.")
}

pub fn main() {
    let args: Vec<String> = std::env::args().collect();
    let n_trials: usize = if let [_, arg1] = &args[..] {
        if let Ok(n) = arg1.parse() {
            n
        } else {
            eprintln! {"The number of trials must be an unsigned integer."};
            std::process::exit(1);
        }
    } else {
        eprintln!("The program must be called with one argument, the number of trials.");
        std::process::exit(1);
    };
    let mut random_bits = RandBoolIter::<rand::rngs::ThreadRng>::default();
    let sum_of_results: u64 = std::iter::repeat_with(|| simulate_one_family(&mut random_bits))
        .take(n_trials)
        .sum();
    let mean = sum_of_results as f64 / n_trials as f64;
    
    print!("The mean of {} trials was {}, ", n_trials, mean);
    const EXPECTED: f64 = 3.0;
    println!("which is within ±{}% of the expected value.",
       f64::abs(100.0*(mean-EXPECTED)/EXPECTED));
}
```
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.