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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
use indexmap::IndexMap;
use ndarray::prelude::*;
use ndarray::{Data, DataMut, Slice};
use rand::prelude::*;
use rand::thread_rng;

/// Methods for sorting and partitioning 1-D arrays.
pub trait Sort1dExt<A, S>
where
    S: Data<Elem = A>,
{
    /// Return the element that would occupy the `i`-th position if
    /// the array were sorted in increasing order.
    ///
    /// The array is shuffled **in place** to retrieve the desired element:
    /// no copy of the array is allocated.
    /// After the shuffling, all elements with an index smaller than `i`
    /// are smaller than the desired element, while all elements with
    /// an index greater or equal than `i` are greater than or equal
    /// to the desired element.
    ///
    /// No other assumptions should be made on the ordering of the
    /// elements after this computation.
    ///
    /// Complexity ([quickselect](https://en.wikipedia.org/wiki/Quickselect)):
    /// - average case: O(`n`);
    /// - worst case: O(`n`^2);
    /// where n is the number of elements in the array.
    ///
    /// **Panics** if `i` is greater than or equal to `n`.
    fn get_from_sorted_mut(&mut self, i: usize) -> A
    where
        A: Ord + Clone,
        S: DataMut;

    /// A bulk version of [`get_from_sorted_mut`], optimized to retrieve multiple
    /// indexes at once.
    /// It returns an `IndexMap`, with indexes as keys and retrieved elements as
    /// values.
    /// The `IndexMap` is sorted with respect to indexes in increasing order:
    /// this ordering is preserved when you iterate over it (using `iter`/`into_iter`).
    ///
    /// **Panics** if any element in `indexes` is greater than or equal to `n`,
    /// where `n` is the length of the array..
    ///
    /// [`get_from_sorted_mut`]: #tymethod.get_from_sorted_mut
    fn get_many_from_sorted_mut<S2>(&mut self, indexes: &ArrayBase<S2, Ix1>) -> IndexMap<usize, A>
    where
        A: Ord + Clone,
        S: DataMut,
        S2: Data<Elem = usize>;

    /// Partitions the array in increasing order based on the value initially
    /// located at `pivot_index` and returns the new index of the value.
    ///
    /// The elements are rearranged in such a way that the value initially
    /// located at `pivot_index` is moved to the position it would be in an
    /// array sorted in increasing order. The return value is the new index of
    /// the value after rearrangement. All elements smaller than the value are
    /// moved to its left and all elements equal or greater than the value are
    /// moved to its right. The ordering of the elements in the two partitions
    /// is undefined.
    ///
    /// `self` is shuffled **in place** to operate the desired partition:
    /// no copy of the array is allocated.
    ///
    /// The method uses Hoare's partition algorithm.
    /// Complexity: O(`n`), where `n` is the number of elements in the array.
    /// Average number of element swaps: n/6 - 1/3 (see
    /// [link](https://cs.stackexchange.com/questions/11458/quicksort-partitioning-hoare-vs-lomuto/11550))
    ///
    /// **Panics** if `pivot_index` is greater than or equal to `n`.
    ///
    /// # Example
    ///
    /// ```
    /// use ndarray::array;
    /// use ndarray_stats::Sort1dExt;
    ///
    /// let mut data = array![3, 1, 4, 5, 2];
    /// let pivot_index = 2;
    /// let pivot_value = data[pivot_index];
    ///
    /// // Partition by the value located at `pivot_index`.
    /// let new_index = data.partition_mut(pivot_index);
    /// // The pivot value is now located at `new_index`.
    /// assert_eq!(data[new_index], pivot_value);
    /// // Elements less than that value are moved to the left.
    /// for i in 0..new_index {
    ///     assert!(data[i] < pivot_value);
    /// }
    /// // Elements greater than or equal to that value are moved to the right.
    /// for i in (new_index + 1)..data.len() {
    ///      assert!(data[i] >= pivot_value);
    /// }
    /// ```
    fn partition_mut(&mut self, pivot_index: usize) -> usize
    where
        A: Ord + Clone,
        S: DataMut;

    private_decl! {}
}

impl<A, S> Sort1dExt<A, S> for ArrayBase<S, Ix1>
where
    S: Data<Elem = A>,
{
    fn get_from_sorted_mut(&mut self, i: usize) -> A
    where
        A: Ord + Clone,
        S: DataMut,
    {
        let n = self.len();
        if n == 1 {
            self[0].clone()
        } else {
            let mut rng = thread_rng();
            let pivot_index = rng.gen_range(0, n);
            let partition_index = self.partition_mut(pivot_index);
            if i < partition_index {
                self.slice_axis_mut(Axis(0), Slice::from(..partition_index))
                    .get_from_sorted_mut(i)
            } else if i == partition_index {
                self[i].clone()
            } else {
                self.slice_axis_mut(Axis(0), Slice::from(partition_index + 1..))
                    .get_from_sorted_mut(i - (partition_index + 1))
            }
        }
    }

    fn get_many_from_sorted_mut<S2>(&mut self, indexes: &ArrayBase<S2, Ix1>) -> IndexMap<usize, A>
    where
        A: Ord + Clone,
        S: DataMut,
        S2: Data<Elem = usize>,
    {
        let mut deduped_indexes: Vec<usize> = indexes.to_vec();
        deduped_indexes.sort_unstable();
        deduped_indexes.dedup();

        get_many_from_sorted_mut_unchecked(self, &deduped_indexes)
    }

    fn partition_mut(&mut self, pivot_index: usize) -> usize
    where
        A: Ord + Clone,
        S: DataMut,
    {
        let pivot_value = self[pivot_index].clone();
        self.swap(pivot_index, 0);
        let n = self.len();
        let mut i = 1;
        let mut j = n - 1;
        loop {
            loop {
                if i > j {
                    break;
                }
                if self[i] >= pivot_value {
                    break;
                }
                i += 1;
            }
            while pivot_value <= self[j] {
                if j == 1 {
                    break;
                }
                j -= 1;
            }
            if i >= j {
                break;
            } else {
                self.swap(i, j);
                i += 1;
                j -= 1;
            }
        }
        self.swap(0, i - 1);
        i - 1
    }

    private_impl! {}
}

/// To retrieve multiple indexes from the sorted array in an optimized fashion,
/// [get_many_from_sorted_mut] first of all sorts and deduplicates the
/// `indexes` vector.
///
/// `get_many_from_sorted_mut_unchecked` does not perform this sorting and
/// deduplication, assuming that the user has already taken care of it.
///
/// Useful when you have to call [get_many_from_sorted_mut] multiple times
/// using the same indexes.
///
/// [get_many_from_sorted_mut]: ../trait.Sort1dExt.html#tymethod.get_many_from_sorted_mut
pub(crate) fn get_many_from_sorted_mut_unchecked<A, S>(
    array: &mut ArrayBase<S, Ix1>,
    indexes: &[usize],
) -> IndexMap<usize, A>
where
    A: Ord + Clone,
    S: DataMut<Elem = A>,
{
    if indexes.is_empty() {
        return IndexMap::new();
    }

    // Since `!indexes.is_empty()` and indexes must be in-bounds, `array` must
    // be non-empty.
    let mut values = vec![array[0].clone(); indexes.len()];
    _get_many_from_sorted_mut_unchecked(array.view_mut(), &mut indexes.to_owned(), &mut values);

    // We convert the vector to a more search-friendly `IndexMap`.
    indexes.iter().cloned().zip(values.into_iter()).collect()
}

/// This is the recursive portion of `get_many_from_sorted_mut_unchecked`.
///
/// `indexes` is the list of indexes to get. `indexes` is mutable so that it
/// can be used as scratch space for this routine; the value of `indexes` after
/// calling this routine should be ignored.
///
/// `values` is a pre-allocated slice to use for writing the output. Its
/// initial element values are ignored.
fn _get_many_from_sorted_mut_unchecked<A>(
    mut array: ArrayViewMut1<'_, A>,
    indexes: &mut [usize],
    values: &mut [A],
) where
    A: Ord + Clone,
{
    let n = array.len();
    debug_assert!(n >= indexes.len()); // because indexes must be unique and in-bounds
    debug_assert_eq!(indexes.len(), values.len());

    if indexes.is_empty() {
        // Nothing to do in this case.
        return;
    }

    // At this point, `n >= 1` since `indexes.len() >= 1`.
    if n == 1 {
        // We can only reach this point if `indexes.len() == 1`, so we only
        // need to assign the single value, and then we're done.
        debug_assert_eq!(indexes.len(), 1);
        values[0] = array[0].clone();
        return;
    }

    // We pick a random pivot index: the corresponding element is the pivot value
    let mut rng = thread_rng();
    let pivot_index = rng.gen_range(0, n);

    // We partition the array with respect to the pivot value.
    // The pivot value moves to `array_partition_index`.
    // Elements strictly smaller than the pivot value have indexes < `array_partition_index`.
    // Elements greater or equal to the pivot value have indexes > `array_partition_index`.
    let array_partition_index = array.partition_mut(pivot_index);

    // We use a divide-and-conquer strategy, splitting the indexes we are
    // searching for (`indexes`) and the corresponding portions of the output
    // slice (`values`) into pieces with respect to `array_partition_index`.
    let (found_exact, index_split) = match indexes.binary_search(&array_partition_index) {
        Ok(index) => (true, index),
        Err(index) => (false, index),
    };
    let (smaller_indexes, other_indexes) = indexes.split_at_mut(index_split);
    let (smaller_values, other_values) = values.split_at_mut(index_split);
    let (bigger_indexes, bigger_values) = if found_exact {
        other_values[0] = array[array_partition_index].clone(); // Write exactly found value.
        (&mut other_indexes[1..], &mut other_values[1..])
    } else {
        (other_indexes, other_values)
    };

    // We search recursively for the values corresponding to strictly smaller
    // indexes to the left of `partition_index`.
    _get_many_from_sorted_mut_unchecked(
        array.slice_axis_mut(Axis(0), Slice::from(..array_partition_index)),
        smaller_indexes,
        smaller_values,
    );

    // We search recursively for the values corresponding to strictly bigger
    // indexes to the right of `partition_index`. Since only the right portion
    // of the array is passed in, the indexes need to be shifted by length of
    // the removed portion.
    bigger_indexes
        .iter_mut()
        .for_each(|x| *x -= array_partition_index + 1);
    _get_many_from_sorted_mut_unchecked(
        array.slice_axis_mut(Axis(0), Slice::from(array_partition_index + 1..)),
        bigger_indexes,
        bigger_values,
    );
}