sliding-quantile 0.1.0

// Based on the C# reference implementation by Andrey Akinshin.
// https://aakinshin.net/posts/partitioning-heaps-quantile-estimator3/
//
// Copyright (c) 2020–2025 Andrey Akinshin
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

mod aliases;

#[cfg(test)]
mod tests;

use core::ops::{ControlFlow, Index, IndexMut};

pub use aliases::*;

/// Moving quantile over a sliding window based on partitioning heaps.
pub struct Quantile<Num, NumArr, UsizeArr>
where
    Num: num_traits::Float,
{
    /// The number of most recent observations we calculate the desired quantile for.
    window_size: usize, // const

    /// The target quantile in `0.0..=1.0`.
    ///
    /// For example, `0.5` (median) or `0.95` (95th percentile).
    ///
    /// Formally, given a random variable `X` and a probability `p`, the quantile function
    /// finds a value `x` so that `P(X <= x) = p`. `X` represents the values we encounter.
    /// Our threshold value `x` is at the root that connects the two heaps.
    probability: Num, // const

    /// Heap index of the root element that represents the target quantile position.
    ///
    /// This is the index that connects the lower heap (elements less than or equal to the root)
    /// and upper heap (elements larger than or equal to the root).
    ///
    /// Let's say `window_size = 10` and `p = 0.5` (median). Remember that `p` percent of
    /// observations are supposed to be less than or equal to the value at this index.
    /// Then it makes sense to choose `root_heap_idx = 4`, because root and lower heap
    /// make up exactly half of the window.
    ///
    /// For the correct proportion of elements on each side, we calculate
    /// `root_heap_idx = floor((window_size - 1) * probability)`.
    root_heap_idx: usize, // const

    /// Stores the `n=window_size` most recent observations.
    ///
    /// Must be ordered to uphold the heap-partitioning property and to place the
    /// root that approximates the targe quantile at `root_heap_idx`.
    heap: NumArr,

    /// Total number of observations.
    ///
    /// When the window fills up, `total_elem_count % window_size` is the index of the
    /// oldest observation. New observations are always inserted at
    /// `heap[total_elem_count % window_size]`.
    total_elem_count: usize,

    /// Ring buffer that maps sliding window elements to heap indices.
    ///
    /// The heap index of new observations is always inserted at `total_elem_count % window_size`.
    /// If the window is full, this is also the index of the oldest observation. Therefore,
    /// this buffer allows us to efficiently replace the oldest observation with a new one.
    elem_to_heap_idx: UsizeArr,

    /// Maps heap indices to sliding window elements.
    ///
    /// When inserting new observations, we usually need to swap heap elements to uphold
    /// its properties. Swapping elements means they swap heap indices. We need this mapping
    /// to know where to swap them in `elem_to_heap_idx`.
    heap_to_elem_idx: UsizeArr,

    /// Number of elements in the lower heap (elements <= than the quantile).
    ///
    /// Once the initial window is filled, `lower_heap_size == root_heap_idx`.
    lower_heap_size: usize,

    /// Number of elements in the upper heap (elements >= than the quantile).
    ///
    /// Once the initial window is filled, `upper_heap_size == window_size - lower_heap_size - 1`.
    upper_heap_size: usize,
}

#[derive(Clone, Copy)]
enum Heap {
    /// The lower heap is a max-heap containing elements smaller than the root.
    ///
    /// It occupies the positions `0..root_heap_idx`. The parent is always greater
    /// than or equal to its children, so larger values bubble up toward the root boundary.
    Lower,

    /// The upper heap is a min-heap containing elements larger than the root.
    ///
    /// It occupies the positions `root_heap_idx + 1..window_size`. The parent is always less
    /// than or equal to its children, so smaller values bubble up toward the root boundary.
    Upper,
}

#[cfg(any(feature = "std", feature = "alloc"))]
impl<Num> Quantile<Num, alloc::vec::Vec<Num>, alloc::vec::Vec<usize>>
where
    Num: num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive,
{
    /// Create a new quantile tracker with a specific window size.
    ///
    /// # Arguments
    /// * `probability` -- Probability of the quantile to compute. For example, `0.5` for the
    ///   median, or `0.75` for the 75th percentile.
    /// * `window_size` -- Sliding window size. The target quantile is tracked for
    ///   this many most recent observations. The algorithm is intended for moderate window sizes
    ///   of roughly 10 to 10,000 elements.
    ///
    /// # Panics
    /// * when `window_size` is zero
    /// * when `probability` is outside `0..=1`
    #[must_use]
    pub fn new(probability: Num, window_size: usize) -> Self {
        assert!(window_size > 0_usize, "zero window size");
        assert!(probability.is_finite(), "probability must be finite");
        assert!(probability >= Num::zero(), "probability must be >= 0");
        assert!(probability <= Num::one(), "probability must be <= 1");

        let root_heap_idx =
            Num::from_usize(window_size - 1).unwrap_or(Num::max_value()) * probability;
        let root_heap_idx = root_heap_idx.floor().to_usize().unwrap();

        Self {
            window_size,
            probability,
            heap: alloc::vec![Num::nan(); window_size],
            heap_to_elem_idx: alloc::vec![0_usize; window_size],
            elem_to_heap_idx: alloc::vec![0_usize; window_size],
            root_heap_idx,
            lower_heap_size: 0_usize,
            upper_heap_size: 0_usize,
            total_elem_count: 0_usize,
        }
    }
}

impl<Num, const WIN_SIZE: usize> Quantile<Num, [Num; WIN_SIZE], [usize; WIN_SIZE]>
where
    Num: num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive,
{
    /// Create a new quantile tracker with a specific window size.
    ///
    /// # Arguments
    /// * `window_size` -- Sliding window size. The target quantile is tracked for
    ///   this many most recent observations. The algorithm is intended for moderate window sizes
    ///   of roughly 10 to 10,000 elements.
    /// * `probability` -- Probability of the quantile to compute. For example, `0.5` for the
    ///   median, or `0.75` for the 75th percentile.
    ///
    /// # Panics
    /// * when `window_size` is zero
    /// * when `probability` is outside `0..=1`
    #[must_use]
    pub fn new(probability: Num) -> Self {
        assert!(WIN_SIZE > 0_usize, "zero window size");
        assert!(probability.is_finite(), "probability must be finite");
        assert!(probability >= Num::zero(), "probability must be >= 0");
        assert!(probability <= Num::one(), "probability must be <= 1");

        let root_heap_idx = Num::from_usize(WIN_SIZE - 1).unwrap_or(Num::max_value()) * probability;
        let root_heap_idx = root_heap_idx.floor().to_usize().unwrap();

        Self {
            window_size: WIN_SIZE,
            probability,
            heap: [Num::nan(); WIN_SIZE],
            heap_to_elem_idx: [0_usize; WIN_SIZE],
            elem_to_heap_idx: [0_usize; WIN_SIZE],
            root_heap_idx,
            lower_heap_size: 0_usize,
            upper_heap_size: 0_usize,
            total_elem_count: 0_usize,
        }
    }
}

impl<Num, NumArr, UsizeArr> Quantile<Num, NumArr, UsizeArr>
where
    Num: num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive,
    NumArr: Index<usize, Output = Num> + IndexMut<usize, Output = Num>,
    UsizeArr: Index<usize, Output = usize> + IndexMut<usize, Output = usize>,
{
    /// Add an observation and update the quantile estimate.
    ///
    /// If the window is full, the oldest observation will be removed.
    ///
    /// # Arguments
    /// * `value` -- the next observation in your data stream
    #[inline]
    pub fn add(&mut self, value: Num) {
        let elem_idx = self.total_elem_count % self.window_size;
        self.total_elem_count += 1;

        // window is full - replace the oldest element and sift for the heap property
        if self.total_elem_count > self.window_size {
            let heap_idx = self.elem_to_heap_idx[elem_idx];
            self.insert(heap_idx, elem_idx, value);
            self.sift(heap_idx);
            return;
        }

        // first element - just insert at the root, no sifting needed
        if self.total_elem_count == 1 {
            core::hint::cold_path();
            self.insert(self.root_heap_idx, elem_idx, value);
            return;
        }

        // window is filling - pick a heap to extend and sift
        let desired_lower_heap_size =
            Num::from_usize(self.total_elem_count - 1).unwrap() * self.probability;
        let desired_lower_heap_size = desired_lower_heap_size.floor().to_usize().unwrap();

        let heap_idx = if self.lower_heap_size < desired_lower_heap_size {
            self.lower_heap_size += 1;
            self.root_heap_idx - self.lower_heap_size
        } else {
            self.upper_heap_size += 1;
            self.root_heap_idx + self.upper_heap_size
        };
        self.insert(heap_idx, elem_idx, value);
        self.sift(heap_idx);
    }

    #[inline(always)]
    fn insert(&mut self, heap_idx: usize, elem_idx: usize, value: Num) {
        self.heap[heap_idx] = value;
        self.heap_to_elem_idx[heap_idx] = elem_idx;
        self.elem_to_heap_idx[elem_idx] = heap_idx;
    }

    /// After inserting a new value, heap properties may be violated.
    /// The sift operation restores these properties.
    ///
    /// > The sift operation is convergent. The sift process eventually finds the correct
    /// > element position regardless of insertion location or initial heap property violations.
    /// > This robustness keeps quantile relationships intact as the window slides.
    fn sift(&mut self, mut heap_idx: usize) {
        while let ControlFlow::Continue(new_heap_idx) = self.sift_once(heap_idx) {
            heap_idx = new_heap_idx;
        }
    }

    #[inline(always)]
    fn sift_once(&mut self, heap_idx: usize) -> ControlFlow<(), usize> {
        // the root can be displaced by elements from either side
        if heap_idx == self.root_heap_idx {
            if self.lower_heap_size > 0 {
                let new_heap_idx =
                    self.swap_with_children(Heap::Lower, heap_idx, Some(heap_idx - 1), None);
                if new_heap_idx != heap_idx {
                    return ControlFlow::Continue(new_heap_idx);
                }
            }

            if self.upper_heap_size > 0 {
                let new_heap_idx =
                    self.swap_with_children(Heap::Upper, heap_idx, Some(heap_idx + 1), None);
                if new_heap_idx != heap_idx {
                    return ControlFlow::Continue(new_heap_idx);
                }
            }

            core::hint::cold_path();
            return ControlFlow::Break(());
        }

        let heap = if heap_idx > self.root_heap_idx {
            Heap::Upper
        } else {
            Heap::Lower
        };
        let heap_parent_idx = match heap {
            Heap::Upper => self.root_heap_idx + (heap_idx - self.root_heap_idx) / 2,
            Heap::Lower => self.root_heap_idx - (self.root_heap_idx - heap_idx) / 2,
        };

        // sift upward: if the newly inserted value violates the heap property with respect to
        // its parent, swap their positions and continue. This handles cases where a large value
        // is inserted into the lower heap or a small value into the upper heap.
        if self.should_swap(heap, heap_parent_idx, heap_idx) {
            self.swap(heap_idx, heap_parent_idx);
            return ControlFlow::Continue(heap_parent_idx);
        }

        let heap_child_idx1 = (2 * heap_idx).checked_sub(self.root_heap_idx);
        let heap_child_idx2 = if heap_idx > self.root_heap_idx {
            heap_child_idx1.map(|i| i + 1)
        } else {
            heap_child_idx1.and_then(|i| i.checked_sub(1))
        };

        // sift downward: compare the current element with its children and swap with the
        // child that best satisfies the heap property. For the lower heap (max-heap),
        // this means swapping with the larger child. For the upper heap (min-heap),
        // it means swapping with the smaller child.
        let new_heap_idx =
            self.swap_with_children(heap, heap_idx, heap_child_idx1, heap_child_idx2);

        if new_heap_idx != heap_idx {
            return ControlFlow::Continue(new_heap_idx);
        }

        ControlFlow::Break(())
    }

    #[inline(always)]
    fn should_swap(&self, heap: Heap, parent_idx: usize, child_idx: usize) -> bool {
        match heap {
            // check if the current node violates the max-heap property (lower heap)
            Heap::Lower => self.heap[parent_idx] < self.heap[child_idx],
            // check if the current node violates the min-heap property (upper heap)
            Heap::Upper => self.heap[parent_idx] > self.heap[child_idx],
        }
    }

    #[inline(always)]
    fn swap(&mut self, heap_idx1: usize, heap_idx2: usize) {
        let elem_idx1 = self.heap_to_elem_idx[heap_idx1];
        let elem_idx2 = self.heap_to_elem_idx[heap_idx2];
        let value1 = self.heap[heap_idx1];
        let value2 = self.heap[heap_idx2];

        // swap the actual values in the heap
        self.heap[heap_idx1] = value2;
        self.heap[heap_idx2] = value1;

        // keep the heap-to-element mapping correct
        self.heap_to_elem_idx[heap_idx1] = elem_idx2;
        self.heap_to_elem_idx[heap_idx2] = elem_idx1;

        // keep the element-to-heap mapping correct
        self.elem_to_heap_idx[elem_idx1] = heap_idx2;
        self.elem_to_heap_idx[elem_idx2] = heap_idx1;
    }

    #[inline(always)]
    fn is_heap_idx(&self, heap_idx: usize) -> bool {
        // check if above or at lower heap limit & below or at upper heap limit
        let min_lower_heap_idx = self.root_heap_idx - self.lower_heap_size;
        let max_upper_heap_idx = self.root_heap_idx + self.upper_heap_size;
        heap_idx >= min_lower_heap_idx && heap_idx <= max_upper_heap_idx
    }

    /// Check if swapping the element with a child is necessary to maintain heap properties.
    /// Return the new index of the element after swaps.
    fn swap_with_children(
        &mut self,
        heap: Heap,
        heap_idx: usize,
        maybe_left_child_idx: Option<usize>,
        maybe_right_child_idx: Option<usize>,
    ) -> usize {
        let maybe_left_child_idx = maybe_left_child_idx.filter(|i| self.is_heap_idx(*i));
        let maybe_right_child_idx = maybe_right_child_idx.filter(|i| self.is_heap_idx(*i));

        match (maybe_left_child_idx, maybe_right_child_idx) {
            (None, None) => {
                // no children, no swap
                heap_idx
            }
            (Some(left_child_idx), None) if self.should_swap(heap, heap_idx, left_child_idx) => {
                // single child and swap is necessary
                self.swap(heap_idx, left_child_idx);
                left_child_idx
            }
            (Some(_), None) => {
                // single child but must not swap
                heap_idx
            }
            (Some(left_child_idx), Some(right_child_idx))
                if self.should_swap(heap, heap_idx, left_child_idx)
                    || self.should_swap(heap, heap_idx, right_child_idx) =>
            {
                // two children and swap is necessary - pick the right one for swapping
                let swap_child_idx = if self.should_swap(heap, left_child_idx, right_child_idx) {
                    // right child should become the parent of left child
                    right_child_idx
                } else {
                    // left child should become the parent of right child
                    left_child_idx
                };
                self.swap(heap_idx, swap_child_idx);
                swap_child_idx
            }
            (Some(_), Some(_)) => {
                // two children but must not swap
                heap_idx
            }
            (None, Some(_)) => {
                // SAFETY: in no case should there be a right child but no left child
                unsafe { core::hint::unreachable_unchecked() }
            }
        }
    }

    /// The target quantile for the last `n ≤ window_size` observations.
    ///
    /// For example, when initialized with `probability = 0.2`, this returns a
    /// value ≥ 20% of the observations in the sliding window, and ≤ 80% of them.
    ///
    /// This returns `NAN` if zero values have been observed so far.
    #[inline]
    #[must_use]
    pub fn get(&self) -> Num {
        if self.total_elem_count == 0 {
            core::hint::cold_path();
            return Num::nan();
        }

        // single observation or no observation > the target quantile
        if self.total_elem_count == 1 || self.upper_heap_size == 0 {
            core::hint::cold_path();
            return self.heap[self.root_heap_idx];
        }

        // Hyndman-Fan Type 7
        // root element represents one boundary of the quantile estimate
        let a = self.heap[self.root_heap_idx];

        // the first element of the upper heap provides the other boundary
        let b = self.heap[self.root_heap_idx + 1];

        // these two elements provide the order statistics needed for linear interpolation

        // for a window of size n and probability p, the theoretical position is:
        //      h = (n - 1) * p
        let n = self.total_elem_count.min(self.window_size);
        let h = Num::from_usize(n - 1).unwrap() * self.probability;

        // if this position falls between two elements at positions floor(h) and ceil(h),
        // the quantile becomes:
        //      q = element[floor(h)] + (h - floor(h)) * (element[ceil(h)] - element[floor(h)])
        a + (h - h.floor()) * (b - a)
    }
}