Oeufman/src/node.rs

use std::collections::HashMap;
use std::collections::VecDeque;

pub struct Node {
    pub val: Option<char>,
    pub l: Option<Box<Node>>,
    pub r: Option<Box<Node>>,
    pub n: i32,
    pub path: String,
}
impl PartialEq for Node {
    fn eq(&self, other: &Self) -> bool {
        self.n == other.n
    }
}

impl Eq for Node {}

impl PartialOrd for Node {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Node {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.n.cmp(&other.n)
    }
}
impl Node {
    pub fn new() -> Self {
        Node {
            val: None,
            l: None,
            r: None,
            n: 0,
            path: String::new(),
        }
    }

    pub fn set_char(&mut self, v: char) {
        self.val = Some(v);
    }

    pub fn set_counter(&mut self, n: i32) {
        self.n = n;
    }

    pub fn insert(&mut self, l: Node, r: Node) {
        self.l = Some(Box::new(l));
        self.r = Some(Box::new(r));
    }
}

pub fn tree_to_map(root: &mut Node, char_to_bin: bool) -> HashMap<String, String> {
    let mut map: HashMap<String, String> = HashMap::new();
    let mut to_do: VecDeque<&mut Node> = VecDeque::new();
    to_do.push_back(root);

    while let Some(int_root) = to_do.pop_front() {
        if let Some(v) = int_root.val {
            // leaf
            if char_to_bin {
                map.insert(v.to_string(), int_root.path.to_string());
            } else {
                map.insert(int_root.path.to_string(), v.to_string());
            }
        } else {
            // internal
            if let Some(l) = int_root.l.as_mut() {
                l.path = format!("{}0", int_root.path);
                to_do.push_back(l);
            }
            if let Some(r) = int_root.r.as_mut() {
                r.path = format!("{}1", int_root.path);
                to_do.push_back(r);
            }
        }
    }

    map
}

pub fn create_tree(mut nodes: Vec<Node>) -> Node {
    while nodes.len() > 1 {
        //nodes.sort_by_key(|node| node.n);
        let mut new_root = Node::new();
        let l = nodes.remove(0);
        let r = nodes.remove(0);
        new_root.set_counter(l.n + r.n);
        new_root.insert(l, r);
        //nodes.push(new_root);
        match nodes.binary_search(&new_root) {
            Ok(_) => {}
            Err(pos) => {
                nodes.insert(pos, new_root);
            }
        }
    }
    nodes.pop().unwrap()
}

pub fn tree_to_string(root: &mut Node) -> (String, Vec<char>) {
    let mut tree = String::new();
    let mut chars: Vec<char> = Vec::new();
    let mut to_do: VecDeque<&mut Node> = VecDeque::new();
    to_do.push_back(root);

    while let Some(n) = to_do.pop_front() {
        match n.val {
            Some(c) => {
                tree.push('0');
                chars.push(c);
            }
            None => {
                tree.push('1');
                if let Some(l) = n.l.as_mut() {
                    to_do.push_back(l);
                }
                if let Some(r) = n.r.as_mut() {
                    to_do.push_back(r);
                }
            }
        }
    }

    (tree, chars)
}

pub fn read_tree(binary: &str, chars: &Vec<char>) -> Node {
    let mut stack_bin = String::from(binary);
    let mut stack_chars = chars.clone();
    let mut to_do: VecDeque<&mut Node> = VecDeque::new();
    let mut root = Node::new();

    if stack_bin.remove(0) == '1' {
        root.path = String::from("");
        to_do.push_back(&mut root);
    } else {
        root.val = Some(stack_chars.remove(0));
        root.path = String::from("0");
        return root;
    }

    while let Some(int_root) = to_do.pop_front() {
        int_root.insert(Node::new(), Node::new());
        if let Some(left) = int_root.l.as_mut() {
            if stack_bin.remove(0) == '1' {
                left.path = String::from("1");
                to_do.push_back(left);
            } else {
                left.path = String::from("0");
                left.val = Some(stack_chars.remove(0));
            }
        }
        if let Some(right) = int_root.r.as_mut() {
            if stack_bin.remove(0) == '1' {
                right.path = String::from("1");
                to_do.push_back(right);
            } else {
                right.path = String::from("0");
                right.val = Some(stack_chars.remove(0));
            }
        }
    }
    root
}
pub fn nodes_at_vec(map: HashMap<char, i32>) -> Vec<Node> {
    let mut nodes = Vec::new();

    for (key, value) in &map {
        let mut node = Node::new();
        node.set_char(*key);
        node.set_counter(*value);
        nodes.push(node);
    }
    nodes.sort();
    nodes
}