linfa_trees/decision_trees/
iter.rs

1use std::collections::VecDeque;
2use std::fmt::Debug;
3use std::iter::Iterator;
4
5use super::TreeNode;
6use linfa::{Float, Label};
7
8/// Level-order (BFT) iterator of nodes in a decision tree
9#[derive(Debug, Clone, PartialEq)]
10pub struct NodeIter<'a, F, L> {
11    queue: VecDeque<&'a TreeNode<F, L>>,
12}
13
14impl<'a, F, L> NodeIter<'a, F, L> {
15    pub fn new(queue: VecDeque<&'a TreeNode<F, L>>) -> Self {
16        NodeIter { queue }
17    }
18}
19
20impl<'a, F: Float, L: Debug + Label> Iterator for NodeIter<'a, F, L> {
21    type Item = &'a TreeNode<F, L>;
22
23    fn next(&mut self) -> Option<Self::Item> {
24        #[allow(clippy::manual_inspect)]
25        self.queue.pop_front().map(|node| {
26            node.children()
27                .into_iter()
28                .filter_map(|x| x.as_ref())
29                .for_each(|child| self.queue.push_back(child));
30            node
31        })
32    }
33}