Advertisement
NLinker

Trees on arena allocators

Jul 27th, 2021
1,641
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Rust 3.44 KB | None | 0 0
  1. // https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=7cccabd269fd1ee8f61ff23fd79117e7
  2.  
  3. #[derive(Debug)]
  4. struct Node<T>
  5. where
  6.     T: PartialEq,
  7. {
  8.     idx: usize,
  9.     val: T,
  10.     parent: Option<usize>,
  11.     children: Vec<usize>,
  12. }
  13.  
  14. impl<T> Node<T>
  15. where
  16.     T: PartialEq,
  17. {
  18.     fn new(idx: usize, val: T) -> Self {
  19.         Self {
  20.             idx,
  21.             val,
  22.             parent: None,
  23.             children: vec![],
  24.         }
  25.     }
  26. }
  27.  
  28. #[derive(Debug, Default)]
  29. struct ArenaTree<T>
  30. where
  31.     T: PartialEq,
  32. {
  33.     arena: Vec<Node<T>>,
  34. }
  35.  
  36. impl<T> ArenaTree<T>
  37. where
  38.     T: PartialEq,
  39. {
  40.     fn node(&mut self, val: T) -> usize {
  41.         //first see if it exists
  42.         for node in &self.arena {
  43.             if node.val == val {
  44.                 return node.idx;
  45.             }
  46.         }
  47.         // Otherwise, add new node
  48.         let idx = self.arena.len();
  49.         self.arena.push(Node::new(idx, val));
  50.         idx
  51.     }
  52.     fn size(&self) -> usize {
  53.         self.arena.len()
  54.     }
  55.     fn edges(&self) -> usize {
  56.         self.arena
  57.             .iter()
  58.             .fold(0, |acc, node| acc + node.children.len())
  59.     }
  60.  
  61.     fn depth(&self, idx: usize) -> usize {
  62.         match self.arena[idx].parent {
  63.             Some(id) => 1 + self.depth(id),
  64.             None => 0,
  65.         }
  66.     }
  67.     fn depth_to_target(&self, idx: usize, target: &T) -> Option<usize> {
  68.         // are we here?  If so, Some(0)
  69.         if target == &self.arena[idx].val {
  70.             return Some(0);
  71.         }
  72.  
  73.         // If not, try all children
  74.         for p in &self.arena[idx].children {
  75.             if let Some(x) = self.depth_to_target(*p, &target) {
  76.                 return Some(1 + x);
  77.             }
  78.         }
  79.         // If it cant be found, return None
  80.         None
  81.     }
  82.     fn distance_between(&mut self, from: T, target: T) -> usize {
  83.         // If it's not in the tree, this will add a new unconnected node
  84.         // the final function will still return None
  85.         let start_node = self.node(from);
  86.         let mut ret = 0;
  87.         // Start traversal
  88.         let mut trav = &self.arena[start_node];
  89.         // Explore all children, then hop up one
  90.         while let Some(inner) = trav.parent {
  91.             if let Some(x) = self.depth_to_target(inner, &target) {
  92.                 ret += x;
  93.                 break;
  94.             }
  95.             trav = &self.arena[inner];
  96.             ret += 1;
  97.         }
  98.         // don't go all the way to target, just orbit
  99.         if ret > 0 {
  100.             ret - 1
  101.         } else {
  102.             ret
  103.         }
  104.     }
  105. }
  106.  
  107. fn main() {
  108.     let mut tree: ArenaTree<String> = ArenaTree::default();
  109.     let hello = tree.node("Hello".into());
  110.     let world = tree.node("World".into());
  111.     tree.arena[hello].children.push(world);
  112.     tree.arena[world].parent = Some(hello);
  113.     let exclamation = tree.node("!".into());
  114.     tree.arena[exclamation].parent = Some(world);
  115.     tree.arena[world].children.push(exclamation);
  116.     let period = tree.node(".".into());
  117.     let period_2 = tree.node(".".into());
  118.     assert_eq!(period, period_2);
  119.     tree.arena[world].children.push(period);
  120.     tree.arena[period_2].parent = Some(world);
  121.  
  122.     println!(
  123.         "Total nodes: {}\nTotal edges: {}\nDepth of '.': {}\nDistance between World and !: {}",
  124.         tree.size(),
  125.         tree.edges(),
  126.         tree.depth(period),
  127.         tree.distance_between("World".into(), "!".into())
  128.     );
  129. }
  130.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement