dtrees_rs/algorithms/common/heuristics/
helpers.rs

1//! Split scores computed from class counts.
2
3/// Shannon entropy (base 2) of a class distribution.
4#[inline]
5pub fn entropy(distribution: &[usize]) -> f64 {
6    let sum: usize = distribution.iter().sum();
7    if sum == 0 {
8        return 0.0;
9    }
10
11    let mut entropy = 0.0;
12    for &count in distribution.iter() {
13        if count > 0 {
14            let probability = count as f64 / sum as f64;
15            entropy -= probability * probability.log2();
16        }
17    }
18    entropy
19}
20
21/// Weighted Gini impurity of the two children of a split.
22#[inline]
23pub fn gini_index(
24    parent_distribution: &[usize],
25    left_distribution: &[usize],
26    right_distribution: &[usize],
27    _parent_entropy: f64,
28) -> f64 {
29    let total_samples = parent_distribution.iter().sum::<usize>();
30    if total_samples == 0 {
31        return 0.0;
32    }
33
34    let total_weight = total_samples as f64;
35    let left_weight = left_distribution.iter().sum::<usize>() as f64;
36    let right_weight = right_distribution.iter().sum::<usize>() as f64;
37
38    let left_impurity = calculate_branch_impurity(left_distribution, left_weight);
39    let right_impurity = calculate_branch_impurity(right_distribution, right_weight);
40
41    ((left_weight * left_impurity) + (right_weight * right_impurity)) / total_weight
42}
43
44#[inline]
45fn calculate_branch_impurity(distribution: &[usize], total: f64) -> f64 {
46    if total < 1.0 {
47        return 0.0;
48    }
49
50    1.0 - distribution
51        .iter()
52        .map(|&count| {
53            let probability = count as f64 / total;
54            probability * probability
55        })
56        .sum::<f64>()
57}
58
59/// Entropy of the parent minus the weighted entropy of the children.
60#[inline]
61pub fn information_gain(
62    parent_distribution: &[usize],
63    left_distribution: &[usize],
64    right_distribution: &[usize],
65    parent_entropy: f64,
66) -> f64 {
67    let total_count = parent_distribution.iter().sum::<usize>() as f64;
68    if total_count < 1.0 {
69        return 0.0;
70    }
71
72    let left_count = left_distribution.iter().sum::<usize>() as f64;
73    let right_count = right_distribution.iter().sum::<usize>() as f64;
74
75    let left_weight = left_count / total_count;
76    let right_weight = right_count / total_count;
77
78    let left_entropy = entropy(left_distribution);
79    let right_entropy = entropy(right_distribution);
80    parent_entropy - (left_weight * left_entropy + right_weight * right_entropy)
81}
82
83/// Weighted entropy of the two children; infinite when a child is empty.
84#[inline]
85pub fn weighted_entropy(
86    parent_distribution: &[usize],
87    left_distribution: &[usize],
88    right_distribution: &[usize],
89    _parent_entropy: f64,
90) -> f64 {
91    let total_count = parent_distribution.iter().sum::<usize>() as f64;
92    if total_count < 1.0 {
93        return f64::INFINITY;
94    }
95
96    let left_count = left_distribution.iter().sum::<usize>() as f64;
97    let right_count = right_distribution.iter().sum::<usize>() as f64;
98
99    let left_weight = left_count / total_count;
100    let right_weight = right_count / total_count;
101
102    if left_weight < f64::EPSILON || right_weight < f64::EPSILON {
103        return f64::INFINITY;
104    }
105
106    let left_entropy = entropy(left_distribution);
107    let right_entropy = entropy(right_distribution);
108
109    left_weight * left_entropy + right_weight * right_entropy
110}
dtrees_rs/algorithms/common/heuristics/helpers.rs

dtrees_rs/algorithms/common/heuristics/
helpers.rs
