thesis-code/src/tree.rs

96 lines
2.3 KiB
Rust

use num_traits::Float;
#[derive(PartialEq, Eq, Debug)]
pub struct Alternative<T: Eq> {
pub name: T
}
#[derive(PartialEq, Debug)]
pub struct Edge<T: Eq, U: Float> {
pub weight: U,
pub destination: Tree<T, U>
}
#[derive(Debug)]
pub enum Vertex<T: Eq, U: Float> {
NonTerminal(Box<Edge<T, U>>, Box<Edge<T, U>>),
Terminal(Alternative<T>),
}
// implement symmetry
impl<T: Eq, U: Float> PartialEq for Vertex<T, U>{
fn eq(&self, other: &Self) -> bool {
match self {
Vertex::NonTerminal(a, b) => {
if let Vertex::NonTerminal(c, d) = other {
(a == c && b == d) || (a == d && b == c)
} else {
false
}
}
Vertex::Terminal(a) => {
if let Vertex::Terminal(b) = other {
a == b
} else {
false
}
}
}
}
}
#[derive(PartialEq, Debug)]
pub struct Tree<T: Eq, U: Float> {
pub root: Vertex<T, U>
}
impl<T: Eq, U: Float> Tree<T, U> {
fn left_edge_weights(&self) -> U {
match &self.root {
Vertex::NonTerminal(a, _) => {
let left = &*a;
left.weight
+ left.destination.left_edge_weights()
+ left.destination.right_edge_weights()
},
Vertex::Terminal(_) => U::zero()
}
}
fn right_edge_weights(&self) -> U {
match &self.root {
Vertex::NonTerminal(_, b) => {
let right = &*b;
right.weight
+ right.destination.left_edge_weights()
+ right.destination.right_edge_weights()
},
Vertex::Terminal(_) => U::zero()
}
}
pub fn choice_probability(&self, alternative: &T) -> U {
match &self.root {
Vertex::Terminal(alt) => {
if alt.name == *alternative { U::one() } else { U::zero() }
},
Vertex::NonTerminal(a, b) => {
let left = &*a;
let right = &*b;
let left_edge_weights = self.left_edge_weights();
let right_edge_weights = self.right_edge_weights();
let left_choice_probability =
left.destination.choice_probability(alternative);
let right_choice_probability =
right.destination.choice_probability(alternative);
(
left_edge_weights * left_choice_probability
+ right_edge_weights * right_choice_probability
)
/(left_edge_weights + right_edge_weights)
}
}
}
}