1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
use pyo3::prelude::*;

/// A heap data structure implemented in Rust and exposed to Python.
#[pyclass]
pub struct Heap {
    data: Vec<i32>,
}

#[pymethods]
impl Heap {
    /// Creates a new empty Heap.
    #[new]
    fn new() -> Self {
        Heap { data: Vec::new() }
    }

    /// Pushes a value onto the heap.
    ///
    /// # Arguments
    ///
    /// * `value` - The value to be added to the heap.
    fn push(&mut self, value: i32) {
        self.data.push(value);
        self.sift_up();
    }

    /// Pops the maximum value off the heap.
    ///
    /// Returns `Some(i32)` if the heap is not empty, otherwise returns `None`.
    fn pop(&mut self) -> Option<i32> {
        if self.data.is_empty() {
            return None;
        }
        let result = self.data.swap_remove(0);
        self.sift_down();
        Some(result)
    }

    /// Converts a vector into a heap.
    ///
    /// # Arguments
    ///
    /// * `data` - A vector of integers to be heapified.
    fn heapify(&mut self, data: Vec<i32>) {
        self.data = data;
        for i in (0..self.data.len() / 2).rev() {
            self.sift_down_from(i);
        }
    }

    /// Sorts the heap and returns a sorted vector.
    ///
    /// This consumes the heap.
    fn heapsort(&mut self) -> Vec<i32> {
        let mut sorted = Vec::new();
        while let Some(value) = self.pop() {
            sorted.push(value);
        }
        sorted
    }

    /// Returns the heap as a list.
    fn as_list(&self) -> Vec<i32> {
        self.data.clone()
    }

    /// Moves the last element up to maintain heap property.
    fn sift_up(&mut self) {
        let mut idx = self.data.len() - 1;
        while idx > 0 {
            let parent_idx = (idx - 1) / 2;
            if self.data[idx] <= self.data[parent_idx] {
                break;
            }
            self.data.swap(idx, parent_idx);
            idx = parent_idx;
        }
    }

    /// Moves the root element down to maintain heap property.
    fn sift_down(&mut self) {
        self.sift_down_from(0);
    }

    /// Sifts down the element at the given index to maintain heap property.
    ///
    /// # Arguments
    ///
    /// * `idx` - The index of the element to sift down.
    fn sift_down_from(&mut self, mut idx: usize) {
        let len = self.data.len();
        loop {
            let left = 2 * idx + 1;
            let right = 2 * idx + 2;
            let mut largest = idx;

            if left < len && self.data[left] > self.data[largest] {
                largest = left;
            }
            if right < len && self.data[right] > self.data[largest] {
                largest = right;
            }
            if largest == idx {
                break;
            }
            self.data.swap(idx, largest);
            idx = largest;
        }
    }
}

/// Heap sorts a vector of integers.
///
/// # Arguments
///
/// * `data` - A vector of integers to be sorted.
///
/// Returns a new vector with the integers sorted in descending order.
#[pyfunction]
pub fn Heapsort(data: Vec<i32>) -> Vec<i32> {
    let mut heap = Heap::new();
    heap.heapify(data);
    heap.heapsort()
}