<commit_before>use distribution::Distribution;
use random::Source;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl Distribution for Categorical {
    type Value = usize;

    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }

    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + (i as f64 - mu).powi(2) * p)
    }

    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate()
                         .fold(0.0, |sum, (i, p)| sum + (i as f64 - mu).powi(3) * p);
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate()
                         .fold(0.0, |sum, (i, p)| sum + (i as f64 - mu).powi(4) * p);
        kurt / sigma2.powi(2) - 3.0
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {<commit_msg>Cosmetic adjustments<commit_after>use distribution::Distribution;
use random::Source;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl Distribution for Categorical {
    type Value = usize;

    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }

    fn var(&self) -> f64 {
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(2) * p)
    }

    fn skewness(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let skew = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(3) * p);
        skew / (var * var.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let kurt = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(4) * p);
        kurt / var.powi(2) - 3.0
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
{'fraction_non_alphanumeric': 0.1725625539257981, 'fraction_numerical': 0.034296807592752374, 'mean_word_length': 1.4277486910994763, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5426831', 'n_tokens_mistral': 1718, 'n_tokens_neox': 1626, 'n_words': 443}