somatize_study/partition/
grouped.rs1use super::dealing::{assemble, checked, grouped_by, heaviest_first, in_order};
4use super::{Fold, Partition, PartitionError};
5use crate::Samples;
6use std::fmt;
7
8#[derive(Debug, Clone, PartialEq, Eq, Hash)]
13pub struct Grouped {
14 pub k: usize,
16}
17
18impl Grouped {
19 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
37fn 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}