pub fn stratified_split<T: Hash + Eq>(
labels: &[T],
test_ratio: f64,
seed: u64,
) -> Result<IndexSplit, DatasetError>Expand description
Split into train and test index lists that preserve each class’s proportion.
Like train_test_split, but this function draws the split within each
class rather than over the whole dataset. A class that holds 10% of the
samples then holds about 10% of the train set and 10% of the test set. This
matters for the imbalanced loaders: sms_spam (13% spam), covtype (its
rarest cover type is under 0.5%), and kddcup99. An unstratified split can
omit a rare class from the test set completely.
Every class with at least two members contributes at least one row to each side. A class with a single member contributes it to the train set.
§Parameters
labels- The per-sample class labels, of any comparable, hashable type (&str,String,u8,char, and so on). This covers every label type this crate produces.test_ratio- Fraction of each class to place in the test set, in0.0..=1.0.seed- Seed for the internal generator. The same seed always yields the same split.
§Returns
IndexSplit- The(train, test)row indices, each in shuffled order. Together they are a permutation of0..labels.len().
§Errors
DatasetError::ValidationError- Returns this whenlabelsis empty, or whentest_ratiois not a finite value in0.0..=1.0.
§Example
use dataset_ml::preprocessing::stratified_split;
// Nine samples of class "a", one of class "b".
let labels = ["a", "a", "a", "a", "a", "a", "a", "a", "a", "b"];
let (train, test) = stratified_split(&labels, 0.5, 7).unwrap();
// The lone "b" cannot be in both sides, so it stays in the train set.
assert!(train.contains(&9));
assert_eq!(train.len() + test.len(), 10);