fn split_cache()

in gbdt-rs/src/decision_tree.rs [691:797]
95 lines of code
21 McCabe index (conditional complexity)

    fn split_cache(
        mut self,
        left_set: &[usize],
        right_set: &[usize],
        cache: &TrainingCache,
    ) -> (Self, Self) {
        // level 0: return empty SubCache
        if cache.cache_level == 0 {
            return (SubCache::get_empty(), SubCache::get_empty());
        }

        // allocate the vectors
        let mut left_ordered_features: Vec<Vec<(usize, ValueType)>> =
            Vec::with_capacity(cache.feature_size);
        let mut right_ordered_features: Vec<Vec<(usize, ValueType)>> =
            Vec::with_capacity(cache.feature_size);
        let mut left_ordered_residual = Vec::with_capacity(left_set.len());
        let mut right_ordered_residual = Vec::with_capacity(right_set.len());
        for _ in 0..cache.feature_size {
            left_ordered_features.push(Vec::with_capacity(left_set.len()));
            right_ordered_features.push(Vec::with_capacity(right_set.len()));
        }

        // compute two boolean vectors
        let mut left_bool = vec![false; cache.sample_size];
        let mut right_bool = vec![false; cache.sample_size];
        for index in left_set.iter() {
            left_bool[*index] = true;
        }
        for index in right_set.iter() {
            right_bool[*index] = true;
        }

        // If lazy is true, compute the ordered_features from the TrainingCache.
        if self.lazy {
            for (feature_index, feature_vec) in cache.ordered_features.iter().enumerate() {
                for pair in feature_vec.iter() {
                    let (index, value) = *pair;
                    if left_bool[index] {
                        left_ordered_features[feature_index].push((index, value));
                        continue;
                    }
                    if right_bool[index] {
                        right_ordered_features[feature_index].push((index, value));
                    }
                }
            }
        } else {
            // If lazy is false, compute the ordered_features from the current SubCache
            for feature_index in 0..self.ordered_features.len() {
                let feature_vec = &mut self.ordered_features[feature_index];
                for pair in feature_vec.iter() {
                    let (index, value) = *pair;
                    if left_bool[index] {
                        left_ordered_features[feature_index].push((index, value));
                        continue;
                    }
                    if right_bool[index] {
                        right_ordered_features[feature_index].push((index, value));
                    }
                }
                feature_vec.clear();
                feature_vec.shrink_to_fit();
            }
            self.ordered_features.clear();
            self.ordered_features.shrink_to_fit();
        }

        // If lazy is true, compute the ordered_residual from the TrainingCache
        if self.lazy {
            for pair in cache.ordered_residual.iter() {
                let (index, value) = *pair;
                if left_bool[index] {
                    left_ordered_residual.push((index, value));
                    continue;
                }
                if right_bool[index] {
                    right_ordered_residual.push((index, value));
                }
            }
        } else {
            // If lazy is false, compute the ordered_residual from the current SubCache
            for pair in self.ordered_residual.into_iter() {
                let (index, value) = pair;
                if left_bool[index] {
                    left_ordered_residual.push((index, value));
                    continue;
                }
                if right_bool[index] {
                    right_ordered_residual.push((index, value));
                }
            }
        }
        // return result
        (
            SubCache {
                ordered_features: left_ordered_features,
                ordered_residual: left_ordered_residual,
                lazy: false,
            },
            SubCache {
                ordered_features: right_ordered_features,
                ordered_residual: right_ordered_residual,
                lazy: false,
            },
        )
    }