Skip to main content

somatize_study/partition/
grouped.rs

1//! `k` folds where all the samples of a group land on the same side.
2
3use super::dealing::{assemble, checked, grouped_by, heaviest_first, in_order};
4use super::{Fold, Partition, PartitionError};
5use crate::Samples;
6use std::fmt;
7
8/// Grouping is a k-fold **over the groups**, with the samples following theirs.
9/// Needs [`in_groups`](Samples::in_groups) and takes no seed: it places the
10/// biggest groups first into whichever fold is emptiest, which is what keeps the
11/// folds comparable when the groups are not.
12#[derive(Debug, Clone, PartialEq, Eq, Hash)]
13pub struct Grouped {
14    /// How many folds.
15    pub k: usize,
16}
17
18impl Grouped {
19    /// The folds, in order. Two samples with the same group never end up on
20    /// opposite sides.
21    pub fn folds(&self, samples: &Samples) -> Result<Vec<Fold>, PartitionError> {
22        let n = samples.n();
23        let groups = samples.groups().ok_or(PartitionError::NeedsGroups)?;
24        checked(self.k, n)?;
25
26        let mut tests = vec![Vec::new(); self.k];
27        let mut carried = vec![0usize; self.k];
28        for (_, indices) in heaviest_first(grouped_by(groups, &in_order(n)), self.k)? {
29            let fold = emptiest(&carried);
30            carried[fold] += indices.len();
31            tests[fold].extend(indices);
32        }
33        Ok(assemble(n, tests))
34    }
35}
36
37/// The fold carrying the fewest samples, the first of them on a tie.
38fn emptiest(carried: &[usize]) -> usize {
39    let mut best = 0;
40    for (fold, &load) in carried.iter().enumerate() {
41        if load < carried[best] {
42            best = fold;
43        }
44    }
45    best
46}
47
48impl fmt::Display for Grouped {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        write!(f, "grouped:{}", self.k)
51    }
52}
53
54impl From<Grouped> for Partition {
55    fn from(cut: Grouped) -> Self {
56        Self::Grouped(cut)
57    }
58}