heap.rs

#[derive(Debug)]
struct Heap {
    values: Vec<i64>,
}

// ヒープの条件
// 頂点のvの親頂点をpとした時,key[p]>=key[v]が成立する
// 木の高さをhとした時,木の深さをh-1以下の部分については完全二分木
// 木の高さをhとした時,木の深さhの部分については,頂点が左詰めされている
impl Heap {
    pub fn new() -> Self {
        Heap { values: vec![] }
    }

    pub fn find_max(&self) -> Option<i64> {
        if self.values.len() == 0 {
            None
        } else {
            Some(self.values[0])
        }
    }

    pub fn insert(&mut self, value: i64) {
        self.values.push(value);
        let mut current_index = self.values.len() - 1;

        while let Some(parent_index) = self.find_parent_index(current_index) {
            let current_value = self.values[current_index];
            let parent_value = self.values[parent_index];
            if parent_value < current_value {
                self.values[current_index] = parent_value;
                self.values[parent_index] = current_value;
                current_index = parent_index;
            } else {
                break;
            }
        }
    }

    pub fn delete_max(&mut self) {
        if self.values.is_empty() {
            return;
        }

        let last_value = self.values[self.values.len() - 1];
        self.values.remove(self.values.len() - 1);
        self.values[0] = last_value;

        let mut current_index = 0;

        loop {
            if let Some(max_index) = match self.find_children_index_opt(current_index) {
                (Some(left_index), Some(right_index)) => {
                    let left = self.values[left_index];
                    let right = self.values[right_index];
                    let max_index = if left > right {
                        left_index
                    } else {
                        right_index
                    };
                    Some(max_index)
                }
                (Some(left_index), _) => Some(left_index),
                _ => None,
            } {
                let max = self.values[max_index];
                let current_value = self.values[current_index];
                if max > current_value {
                    self.values[current_index] = max;
                    self.values[max_index] = current_value;
                    current_index = max_index;
                } else {
                    return;
                }
            } else {
                return;
            }
        }
    }

    fn find_parent_index(&self, index: usize) -> Option<usize> {
        if index == 0 {
            None
        } else {
            Some((index - 1) / 2)
        }
    }

    fn find_children_index_opt(&self, index: usize) -> (Option<usize>, Option<usize>) {
        if self.values.len() == 0 {
            return (None, None);
        }
        let max_index = self.values.len() - 1;
        let left_index = 2 * index + 1;

        let right_index = 2 * index + 2;
        if right_index <= max_index {
            return (Some(left_index), Some(right_index));
        } else if left_index <= max_index {
            return (Some(left_index), None);
        } else {
            return (None, None);
        }
    }
}

fn main() {
    let mut heap = Heap::new();
    heap.insert(1);
    heap.insert(5);
    heap.insert(3);
    heap.insert(8);
    heap.insert(7);
    heap.insert(16);
    heap.insert(4);
    heap.insert(9);
    println!("{:?}", heap);

    heap.delete_max();
    heap.delete_max();

    println!("{:?}", heap);
    println!("{:?}", heap.find_max());
}