    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
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137

    //! N-best selection.

    use core::iter::Iterator;
    use core::ops::{Deref, DerefMut};

    /// N-best elements.
    pub struct NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
    {
        n: usize,
        f: F,
        candidates: Vec<T>,
    }

    impl<T, K, F> NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
    {
        /// Creates a new [`NBestByKey`].
        pub fn new(n: usize, f: F) -> Self {
            Self {
                n,
                f,
                candidates: Vec::with_capacity(n),
            }
        }

        /// Consumes the [`NBestByKey`] and returns the underlying vector.
        pub fn into_vec(self) -> Vec<T> {
            self.candidates
        }
    }

    impl<T, K, F> NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
        K: PartialOrd,
    {
        /// Pushes a new candidate to the n-best.
        ///
        /// If there are less than `n` candidates, the candidate is pushed to the
        /// current candidates.
        ///
        /// Replaces an item in the current n-best with `candidate`, such that
        /// `candidate` is less than that item (in terms of `F(&T) -> K`).
        /// That replaced item further replaces another item in the current n-best,
        /// such that the replaced item is less than that item.
        /// Repeats the above process until no replacement occurs.
        ///
        /// Returns the item pushed out unless there are less than `n` candidates.
        pub fn push(&mut self, mut candidate: T) -> Option<T> {
            if self.candidates.len() < self.n {
                self.candidates.push(candidate);
                return None;
            }
            while let Some(to_replace) = self.candidates
                .iter_mut()
                .find(|item| (self.f)(&candidate).lt(&(self.f)(item)))
            {
                std::mem::swap(&mut candidate, to_replace);
            }
            Some(candidate)
        }
    }

    impl<T, K, F> Deref for NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
    {
        type Target = [T];

        fn deref(&self) -> &Self::Target {
            &self.candidates
        }
    }

    impl<T, K, F> DerefMut for NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
    {
        fn deref_mut(&mut self) -> &mut Self::Target {
            &mut self.candidates
        }
    }

    impl<T, K, F> Into<Vec<T>> for NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
    {
        fn into(self) -> Vec<T> {
            self.into_vec()
        }
    }

    impl<T, K, F> IntoIterator for NBestByKey<T, K, F>
    where
        F: FnMut(&T) -> K,
    {
        type Item = T;
        type IntoIter = std::vec::IntoIter<Self::Item>;

        fn into_iter(self) -> Self::IntoIter {
            self.candidates.into_iter()
        }
    }

    /// Trait to select n-best elements in a collection.
    ///
    /// Intended to be implemented for [`Iterator`].
    pub trait TakeNBestByKey<T>
    where
        Self: Sized,
    {
        /// Selects n-best elements in the iterator.
        fn n_best_by_key<K, F>(self, n: usize, f: F) -> NBestByKey<T, K, F>
        where
            F: FnMut(&T) -> K,
            K: PartialOrd;
    }

    impl<I, T> TakeNBestByKey<T> for I
    where
        I: Iterator<Item = T> + Sized
    {
        fn n_best_by_key<K, F>(self, n: usize, f: F) -> NBestByKey<T, K, F>
        where
            F: FnMut(&T) -> K,
            K: PartialOrd,
        {
            let mut n_best = NBestByKey::new(n, f);
            for item in self {
                n_best.push(item);
            }
            n_best
        }
    }