Skip to content

Instantly share code, notes, and snippets.

@rust-play
Created December 18, 2019 23:40
Show Gist options
  • Select an option

  • Save rust-play/7cccabd269fd1ee8f61ff23fd79117e7 to your computer and use it in GitHub Desktop.

Select an option

Save rust-play/7cccabd269fd1ee8f61ff23fd79117e7 to your computer and use it in GitHub Desktop.
Code shared from the Rust Playground
#[derive(Debug)]
struct Node<T>
where
T: PartialEq,
{
idx: usize,
val: T,
parent: Option<usize>,
children: Vec<usize>,
}
impl<T> Node<T>
where
T: PartialEq,
{
fn new(idx: usize, val: T) -> Self {
Self {
idx,
val,
parent: None,
children: vec![],
}
}
}
#[derive(Debug, Default)]
struct ArenaTree<T>
where
T: PartialEq,
{
arena: Vec<Node<T>>,
}
impl<T> ArenaTree<T>
where
T: PartialEq,
{
fn node(&mut self, val: T) -> usize {
//first see if it exists
for node in &self.arena {
if node.val == val {
return node.idx;
}
}
// Otherwise, add new node
let idx = self.arena.len();
self.arena.push(Node::new(idx, val));
idx
}
fn size(&self) -> usize {
self.arena.len()
}
fn edges(&self) -> usize {
self.arena
.iter()
.fold(0, |acc, node| acc + node.children.len())
}
fn depth(&self, idx: usize) -> usize {
match self.arena[idx].parent {
Some(id) => 1 + self.depth(id),
None => 0,
}
}
fn depth_to_target(&self, idx: usize, target: &T) -> Option<usize> {
// are we here? If so, Some(0)
if target == &self.arena[idx].val {
return Some(0);
}
// If not, try all children
for p in &self.arena[idx].children {
if let Some(x) = self.depth_to_target(*p, &target) {
return Some(1 + x);
}
}
// If it cant be found, return None
None
}
fn distance_between(&mut self, from: T, target: T) -> usize {
// If it's not in the tree, this will add a new unconnected node
// the final function will still return None
let start_node = self.node(from);
let mut ret = 0;
// Start traversal
let mut trav = &self.arena[start_node];
// Explore all children, then hop up one
while let Some(inner) = trav.parent {
if let Some(x) = self.depth_to_target(inner, &target) {
ret += x;
break;
}
trav = &self.arena[inner];
ret += 1;
}
// don't go all the way to target, just orbit
if ret > 0 {
ret - 1
} else {
ret
}
}
}
fn main() {
let mut tree: ArenaTree<String> = ArenaTree::default();
let hello = tree.node("Hello".into());
let world = tree.node("World".into());
tree.arena[hello].children.push(world);
tree.arena[world].parent = Some(hello);
let exclamation = tree.node("!".into());
tree.arena[exclamation].parent = Some(world);
tree.arena[world].children.push(exclamation);
let period = tree.node(".".into());
let period_2 = tree.node(".".into());
assert_eq!(period, period_2);
tree.arena[world].children.push(period);
tree.arena[period_2].parent = Some(world);
println!(
"Total nodes: {}\nTotal edges: {}\nDepth of '.': {}\nDistance between World and !: {}",
tree.size(),
tree.edges(),
tree.depth(period),
tree.distance_between("World".into(), "!".into())
);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment