rand/seq/
slice.rs

1// Copyright 2018-2023 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! `IndexedRandom`, `IndexedMutRandom`, `SliceRandom`
10
11use super::increasing_uniform::IncreasingUniform;
12use super::index;
13#[cfg(feature = "alloc")]
14use crate::distr::uniform::{SampleBorrow, SampleUniform};
15#[cfg(feature = "alloc")]
16use crate::distr::weighted::{Error as WeightError, Weight};
17use crate::Rng;
18use core::ops::{Index, IndexMut};
19
20/// Extension trait on indexable lists, providing random sampling methods.
21///
22/// This trait is implemented on `[T]` slice types. Other types supporting
23/// [`std::ops::Index<usize>`] may implement this (only [`Self::len`] must be
24/// specified).
25pub trait IndexedRandom: Index<usize> {
26    /// The length
27    fn len(&self) -> usize;
28
29    /// True when the length is zero
30    #[inline]
31    fn is_empty(&self) -> bool {
32        self.len() == 0
33    }
34
35    /// Uniformly sample one element
36    ///
37    /// Returns a reference to one uniformly-sampled random element of
38    /// the slice, or `None` if the slice is empty.
39    ///
40    /// For slices, complexity is `O(1)`.
41    ///
42    /// # Example
43    ///
44    /// ```
45    /// use rand::seq::IndexedRandom;
46    ///
47    /// let choices = [1, 2, 4, 8, 16, 32];
48    /// let mut rng = rand::rng();
49    /// println!("{:?}", choices.choose(&mut rng));
50    /// assert_eq!(choices[..0].choose(&mut rng), None);
51    /// ```
52    fn choose<R>(&self, rng: &mut R) -> Option<&Self::Output>
53    where
54        R: Rng + ?Sized,
55    {
56        if self.is_empty() {
57            None
58        } else {
59            Some(&self[rng.random_range(..self.len())])
60        }
61    }
62
63    /// Uniformly sample `amount` distinct elements from self
64    ///
65    /// Chooses `amount` elements from the slice at random, without repetition,
66    /// and in random order. The returned iterator is appropriate both for
67    /// collection into a `Vec` and filling an existing buffer (see example).
68    ///
69    /// In case this API is not sufficiently flexible, use [`index::sample`].
70    ///
71    /// For slices, complexity is the same as [`index::sample`].
72    ///
73    /// # Example
74    /// ```
75    /// use rand::seq::IndexedRandom;
76    ///
77    /// let mut rng = &mut rand::rng();
78    /// let sample = "Hello, audience!".as_bytes();
79    ///
80    /// // collect the results into a vector:
81    /// let v: Vec<u8> = sample.choose_multiple(&mut rng, 3).cloned().collect();
82    ///
83    /// // store in a buffer:
84    /// let mut buf = [0u8; 5];
85    /// for (b, slot) in sample.choose_multiple(&mut rng, buf.len()).zip(buf.iter_mut()) {
86    ///     *slot = *b;
87    /// }
88    /// ```
89    #[cfg(feature = "alloc")]
90    fn choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Output>
91    where
92        Self::Output: Sized,
93        R: Rng + ?Sized,
94    {
95        let amount = core::cmp::min(amount, self.len());
96        SliceChooseIter {
97            slice: self,
98            _phantom: Default::default(),
99            indices: index::sample(rng, self.len(), amount).into_iter(),
100        }
101    }
102
103    /// Uniformly sample a fixed-size array of distinct elements from self
104    ///
105    /// Chooses `N` elements from the slice at random, without repetition,
106    /// and in random order.
107    ///
108    /// For slices, complexity is the same as [`index::sample_array`].
109    ///
110    /// # Example
111    /// ```
112    /// use rand::seq::IndexedRandom;
113    ///
114    /// let mut rng = &mut rand::rng();
115    /// let sample = "Hello, audience!".as_bytes();
116    ///
117    /// let a: [u8; 3] = sample.choose_multiple_array(&mut rng).unwrap();
118    /// ```
119    fn choose_multiple_array<R, const N: usize>(&self, rng: &mut R) -> Option<[Self::Output; N]>
120    where
121        Self::Output: Clone + Sized,
122        R: Rng + ?Sized,
123    {
124        let indices = index::sample_array(rng, self.len())?;
125        Some(indices.map(|index| self[index].clone()))
126    }
127
128    /// Biased sampling for one element
129    ///
130    /// Returns a reference to one element of the slice, sampled according
131    /// to the provided weights. Returns `None` only if the slice is empty.
132    ///
133    /// The specified function `weight` maps each item `x` to a relative
134    /// likelihood `weight(x)`. The probability of each item being selected is
135    /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`.
136    ///
137    /// For slices of length `n`, complexity is `O(n)`.
138    /// For more information about the underlying algorithm,
139    /// see the [`WeightedIndex`] distribution.
140    ///
141    /// See also [`choose_weighted_mut`].
142    ///
143    /// # Example
144    ///
145    /// ```
146    /// use rand::prelude::*;
147    ///
148    /// let choices = [('a', 2), ('b', 1), ('c', 1), ('d', 0)];
149    /// let mut rng = rand::rng();
150    /// // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c',
151    /// // and 'd' will never be printed
152    /// println!("{:?}", choices.choose_weighted(&mut rng, |item| item.1).unwrap().0);
153    /// ```
154    /// [`choose`]: IndexedRandom::choose
155    /// [`choose_weighted_mut`]: IndexedMutRandom::choose_weighted_mut
156    /// [`WeightedIndex`]: crate::distr::weighted::WeightedIndex
157    #[cfg(feature = "alloc")]
158    fn choose_weighted<R, F, B, X>(
159        &self,
160        rng: &mut R,
161        weight: F,
162    ) -> Result<&Self::Output, WeightError>
163    where
164        R: Rng + ?Sized,
165        F: Fn(&Self::Output) -> B,
166        B: SampleBorrow<X>,
167        X: SampleUniform + Weight + PartialOrd<X>,
168    {
169        use crate::distr::{weighted::WeightedIndex, Distribution};
170        let distr = WeightedIndex::new((0..self.len()).map(|idx| weight(&self[idx])))?;
171        Ok(&self[distr.sample(rng)])
172    }
173
174    /// Biased sampling of `amount` distinct elements
175    ///
176    /// Similar to [`choose_multiple`], but where the likelihood of each element's
177    /// inclusion in the output may be specified. The elements are returned in an
178    /// arbitrary, unspecified order.
179    ///
180    /// The specified function `weight` maps each item `x` to a relative
181    /// likelihood `weight(x)`. The probability of each item being selected is
182    /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`.
183    ///
184    /// If all of the weights are equal, even if they are all zero, each element has
185    /// an equal likelihood of being selected.
186    ///
187    /// This implementation uses `O(length + amount)` space and `O(length)` time
188    /// if the "nightly" feature is enabled, or `O(length)` space and
189    /// `O(length + amount * log length)` time otherwise.
190    ///
191    /// # Known issues
192    ///
193    /// The algorithm currently used to implement this method loses accuracy
194    /// when small values are used for weights.
195    /// See [#1476](https://github.com/rust-random/rand/issues/1476).
196    ///
197    /// # Example
198    ///
199    /// ```
200    /// use rand::prelude::*;
201    ///
202    /// let choices = [('a', 2), ('b', 1), ('c', 1)];
203    /// let mut rng = rand::rng();
204    /// // First Draw * Second Draw = total odds
205    /// // -----------------------
206    /// // (50% * 50%) + (25% * 67%) = 41.7% chance that the output is `['a', 'b']` in some order.
207    /// // (50% * 50%) + (25% * 67%) = 41.7% chance that the output is `['a', 'c']` in some order.
208    /// // (25% * 33%) + (25% * 33%) = 16.6% chance that the output is `['b', 'c']` in some order.
209    /// println!("{:?}", choices.choose_multiple_weighted(&mut rng, 2, |item| item.1).unwrap().collect::<Vec<_>>());
210    /// ```
211    /// [`choose_multiple`]: IndexedRandom::choose_multiple
212    // Note: this is feature-gated on std due to usage of f64::powf.
213    // If necessary, we may use alloc+libm as an alternative (see PR #1089).
214    #[cfg(feature = "std")]
215    fn choose_multiple_weighted<R, F, X>(
216        &self,
217        rng: &mut R,
218        amount: usize,
219        weight: F,
220    ) -> Result<SliceChooseIter<Self, Self::Output>, WeightError>
221    where
222        Self::Output: Sized,
223        R: Rng + ?Sized,
224        F: Fn(&Self::Output) -> X,
225        X: Into<f64>,
226    {
227        let amount = core::cmp::min(amount, self.len());
228        Ok(SliceChooseIter {
229            slice: self,
230            _phantom: Default::default(),
231            indices: index::sample_weighted(
232                rng,
233                self.len(),
234                |idx| weight(&self[idx]).into(),
235                amount,
236            )?
237            .into_iter(),
238        })
239    }
240}
241
242/// Extension trait on indexable lists, providing random sampling methods.
243///
244/// This trait is implemented automatically for every type implementing
245/// [`IndexedRandom`] and [`std::ops::IndexMut<usize>`].
246pub trait IndexedMutRandom: IndexedRandom + IndexMut<usize> {
247    /// Uniformly sample one element (mut)
248    ///
249    /// Returns a mutable reference to one uniformly-sampled random element of
250    /// the slice, or `None` if the slice is empty.
251    ///
252    /// For slices, complexity is `O(1)`.
253    fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Output>
254    where
255        R: Rng + ?Sized,
256    {
257        if self.is_empty() {
258            None
259        } else {
260            let len = self.len();
261            Some(&mut self[rng.random_range(..len)])
262        }
263    }
264
265    /// Biased sampling for one element (mut)
266    ///
267    /// Returns a mutable reference to one element of the slice, sampled according
268    /// to the provided weights. Returns `None` only if the slice is empty.
269    ///
270    /// The specified function `weight` maps each item `x` to a relative
271    /// likelihood `weight(x)`. The probability of each item being selected is
272    /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`.
273    ///
274    /// For slices of length `n`, complexity is `O(n)`.
275    /// For more information about the underlying algorithm,
276    /// see the [`WeightedIndex`] distribution.
277    ///
278    /// See also [`choose_weighted`].
279    ///
280    /// [`choose_mut`]: IndexedMutRandom::choose_mut
281    /// [`choose_weighted`]: IndexedRandom::choose_weighted
282    /// [`WeightedIndex`]: crate::distr::weighted::WeightedIndex
283    #[cfg(feature = "alloc")]
284    fn choose_weighted_mut<R, F, B, X>(
285        &mut self,
286        rng: &mut R,
287        weight: F,
288    ) -> Result<&mut Self::Output, WeightError>
289    where
290        R: Rng + ?Sized,
291        F: Fn(&Self::Output) -> B,
292        B: SampleBorrow<X>,
293        X: SampleUniform + Weight + PartialOrd<X>,
294    {
295        use crate::distr::{weighted::WeightedIndex, Distribution};
296        let distr = WeightedIndex::new((0..self.len()).map(|idx| weight(&self[idx])))?;
297        let index = distr.sample(rng);
298        Ok(&mut self[index])
299    }
300}
301
302/// Extension trait on slices, providing shuffling methods.
303///
304/// This trait is implemented on all `[T]` slice types, providing several
305/// methods for choosing and shuffling elements. You must `use` this trait:
306///
307/// ```
308/// use rand::seq::SliceRandom;
309///
310/// let mut rng = rand::rng();
311/// let mut bytes = "Hello, random!".to_string().into_bytes();
312/// bytes.shuffle(&mut rng);
313/// let str = String::from_utf8(bytes).unwrap();
314/// println!("{}", str);
315/// ```
316/// Example output (non-deterministic):
317/// ```none
318/// l,nmroHado !le
319/// ```
320pub trait SliceRandom: IndexedMutRandom {
321    /// Shuffle a mutable slice in place.
322    ///
323    /// For slices of length `n`, complexity is `O(n)`.
324    /// The resulting permutation is picked uniformly from the set of all possible permutations.
325    ///
326    /// # Example
327    ///
328    /// ```
329    /// use rand::seq::SliceRandom;
330    ///
331    /// let mut rng = rand::rng();
332    /// let mut y = [1, 2, 3, 4, 5];
333    /// println!("Unshuffled: {:?}", y);
334    /// y.shuffle(&mut rng);
335    /// println!("Shuffled:   {:?}", y);
336    /// ```
337    fn shuffle<R>(&mut self, rng: &mut R)
338    where
339        R: Rng + ?Sized;
340
341    /// Shuffle a slice in place, but exit early.
342    ///
343    /// Returns two mutable slices from the source slice. The first contains
344    /// `amount` elements randomly permuted. The second has the remaining
345    /// elements that are not fully shuffled.
346    ///
347    /// This is an efficient method to select `amount` elements at random from
348    /// the slice, provided the slice may be mutated.
349    ///
350    /// If you only need to choose elements randomly and `amount > self.len()/2`
351    /// then you may improve performance by taking
352    /// `amount = self.len() - amount` and using only the second slice.
353    ///
354    /// If `amount` is greater than the number of elements in the slice, this
355    /// will perform a full shuffle.
356    ///
357    /// For slices, complexity is `O(m)` where `m = amount`.
358    fn partial_shuffle<R>(
359        &mut self,
360        rng: &mut R,
361        amount: usize,
362    ) -> (&mut [Self::Output], &mut [Self::Output])
363    where
364        Self::Output: Sized,
365        R: Rng + ?Sized;
366}
367
368impl<T> IndexedRandom for [T] {
369    fn len(&self) -> usize {
370        self.len()
371    }
372}
373
374impl<IR: IndexedRandom + IndexMut<usize> + ?Sized> IndexedMutRandom for IR {}
375
376impl<T> SliceRandom for [T] {
377    fn shuffle<R>(&mut self, rng: &mut R)
378    where
379        R: Rng + ?Sized,
380    {
381        if self.len() <= 1 {
382            // There is no need to shuffle an empty or single element slice
383            return;
384        }
385        self.partial_shuffle(rng, self.len());
386    }
387
388    fn partial_shuffle<R>(&mut self, rng: &mut R, amount: usize) -> (&mut [T], &mut [T])
389    where
390        R: Rng + ?Sized,
391    {
392        let m = self.len().saturating_sub(amount);
393
394        // The algorithm below is based on Durstenfeld's algorithm for the
395        // [Fisher–Yates shuffle](https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm)
396        // for an unbiased permutation.
397        // It ensures that the last `amount` elements of the slice
398        // are randomly selected from the whole slice.
399
400        // `IncreasingUniform::next_index()` is faster than `Rng::random_range`
401        // but only works for 32 bit integers
402        // So we must use the slow method if the slice is longer than that.
403        if self.len() < (u32::MAX as usize) {
404            let mut chooser = IncreasingUniform::new(rng, m as u32);
405            for i in m..self.len() {
406                let index = chooser.next_index();
407                self.swap(i, index);
408            }
409        } else {
410            for i in m..self.len() {
411                let index = rng.random_range(..i + 1);
412                self.swap(i, index);
413            }
414        }
415        let r = self.split_at_mut(m);
416        (r.1, r.0)
417    }
418}
419
420/// An iterator over multiple slice elements.
421///
422/// This struct is created by
423/// [`IndexedRandom::choose_multiple`](trait.IndexedRandom.html#tymethod.choose_multiple).
424#[cfg(feature = "alloc")]
425#[derive(Debug)]
426pub struct SliceChooseIter<'a, S: ?Sized + 'a, T: 'a> {
427    slice: &'a S,
428    _phantom: core::marker::PhantomData<T>,
429    indices: index::IndexVecIntoIter,
430}
431
432#[cfg(feature = "alloc")]
433impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> Iterator for SliceChooseIter<'a, S, T> {
434    type Item = &'a T;
435
436    fn next(&mut self) -> Option<Self::Item> {
437        // TODO: investigate using SliceIndex::get_unchecked when stable
438        self.indices.next().map(|i| &self.slice[i])
439    }
440
441    fn size_hint(&self) -> (usize, Option<usize>) {
442        (self.indices.len(), Some(self.indices.len()))
443    }
444}
445
446#[cfg(feature = "alloc")]
447impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> ExactSizeIterator
448    for SliceChooseIter<'a, S, T>
449{
450    fn len(&self) -> usize {
451        self.indices.len()
452    }
453}
454
455#[cfg(test)]
456mod test {
457    use super::*;
458    #[cfg(feature = "alloc")]
459    use alloc::vec::Vec;
460
461    #[test]
462    fn test_slice_choose() {
463        let mut r = crate::test::rng(107);
464        let chars = [
465            'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
466        ];
467        let mut chosen = [0i32; 14];
468        // The below all use a binomial distribution with n=1000, p=1/14.
469        // binocdf(40, 1000, 1/14) ~= 2e-5; 1-binocdf(106, ..) ~= 2e-5
470        for _ in 0..1000 {
471            let picked = *chars.choose(&mut r).unwrap();
472            chosen[(picked as usize) - ('a' as usize)] += 1;
473        }
474        for count in chosen.iter() {
475            assert!(40 < *count && *count < 106);
476        }
477
478        chosen.iter_mut().for_each(|x| *x = 0);
479        for _ in 0..1000 {
480            *chosen.choose_mut(&mut r).unwrap() += 1;
481        }
482        for count in chosen.iter() {
483            assert!(40 < *count && *count < 106);
484        }
485
486        let mut v: [isize; 0] = [];
487        assert_eq!(v.choose(&mut r), None);
488        assert_eq!(v.choose_mut(&mut r), None);
489    }
490
491    #[test]
492    fn value_stability_slice() {
493        let mut r = crate::test::rng(413);
494        let chars = [
495            'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
496        ];
497        let mut nums = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
498
499        assert_eq!(chars.choose(&mut r), Some(&'l'));
500        assert_eq!(nums.choose_mut(&mut r), Some(&mut 3));
501
502        assert_eq!(
503            &chars.choose_multiple_array(&mut r),
504            &Some(['f', 'i', 'd', 'b', 'c', 'm', 'j', 'k'])
505        );
506
507        #[cfg(feature = "alloc")]
508        assert_eq!(
509            &chars
510                .choose_multiple(&mut r, 8)
511                .cloned()
512                .collect::<Vec<char>>(),
513            &['h', 'm', 'd', 'b', 'c', 'e', 'n', 'f']
514        );
515
516        #[cfg(feature = "alloc")]
517        assert_eq!(chars.choose_weighted(&mut r, |_| 1), Ok(&'i'));
518        #[cfg(feature = "alloc")]
519        assert_eq!(nums.choose_weighted_mut(&mut r, |_| 1), Ok(&mut 2));
520
521        let mut r = crate::test::rng(414);
522        nums.shuffle(&mut r);
523        assert_eq!(nums, [5, 11, 0, 8, 7, 12, 6, 4, 9, 3, 1, 2, 10]);
524        nums = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
525        let res = nums.partial_shuffle(&mut r, 6);
526        assert_eq!(res.0, &mut [7, 12, 6, 8, 1, 9]);
527        assert_eq!(res.1, &mut [0, 11, 2, 3, 4, 5, 10]);
528    }
529
530    #[test]
531    #[cfg_attr(miri, ignore)] // Miri is too slow
532    fn test_shuffle() {
533        let mut r = crate::test::rng(108);
534        let empty: &mut [isize] = &mut [];
535        empty.shuffle(&mut r);
536        let mut one = [1];
537        one.shuffle(&mut r);
538        let b: &[_] = &[1];
539        assert_eq!(one, b);
540
541        let mut two = [1, 2];
542        two.shuffle(&mut r);
543        assert!(two == [1, 2] || two == [2, 1]);
544
545        fn move_last(slice: &mut [usize], pos: usize) {
546            // use slice[pos..].rotate_left(1); once we can use that
547            let last_val = slice[pos];
548            for i in pos..slice.len() - 1 {
549                slice[i] = slice[i + 1];
550            }
551            *slice.last_mut().unwrap() = last_val;
552        }
553        let mut counts = [0i32; 24];
554        for _ in 0..10000 {
555            let mut arr: [usize; 4] = [0, 1, 2, 3];
556            arr.shuffle(&mut r);
557            let mut permutation = 0usize;
558            let mut pos_value = counts.len();
559            for i in 0..4 {
560                pos_value /= 4 - i;
561                let pos = arr.iter().position(|&x| x == i).unwrap();
562                assert!(pos < (4 - i));
563                permutation += pos * pos_value;
564                move_last(&mut arr, pos);
565                assert_eq!(arr[3], i);
566            }
567            for (i, &a) in arr.iter().enumerate() {
568                assert_eq!(a, i);
569            }
570            counts[permutation] += 1;
571        }
572        for count in counts.iter() {
573            // Binomial(10000, 1/24) with average 416.667
574            // Octave: binocdf(n, 10000, 1/24)
575            // 99.9% chance samples lie within this range:
576            assert!(352 <= *count && *count <= 483, "count: {}", count);
577        }
578    }
579
580    #[test]
581    fn test_partial_shuffle() {
582        let mut r = crate::test::rng(118);
583
584        let mut empty: [u32; 0] = [];
585        let res = empty.partial_shuffle(&mut r, 10);
586        assert_eq!((res.0.len(), res.1.len()), (0, 0));
587
588        let mut v = [1, 2, 3, 4, 5];
589        let res = v.partial_shuffle(&mut r, 2);
590        assert_eq!((res.0.len(), res.1.len()), (2, 3));
591        assert!(res.0[0] != res.0[1]);
592        // First elements are only modified if selected, so at least one isn't modified:
593        assert!(res.1[0] == 1 || res.1[1] == 2 || res.1[2] == 3);
594    }
595
596    #[test]
597    #[cfg(feature = "alloc")]
598    #[cfg_attr(miri, ignore)] // Miri is too slow
599    fn test_weighted() {
600        let mut r = crate::test::rng(406);
601        const N_REPS: u32 = 3000;
602        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
603        let total_weight = weights.iter().sum::<u32>() as f32;
604
605        let verify = |result: [i32; 14]| {
606            for (i, count) in result.iter().enumerate() {
607                let exp = (weights[i] * N_REPS) as f32 / total_weight;
608                let mut err = (*count as f32 - exp).abs();
609                if err != 0.0 {
610                    err /= exp;
611                }
612                assert!(err <= 0.25);
613            }
614        };
615
616        // choose_weighted
617        fn get_weight<T>(item: &(u32, T)) -> u32 {
618            item.0
619        }
620        let mut chosen = [0i32; 14];
621        let mut items = [(0u32, 0usize); 14]; // (weight, index)
622        for (i, item) in items.iter_mut().enumerate() {
623            *item = (weights[i], i);
624        }
625        for _ in 0..N_REPS {
626            let item = items.choose_weighted(&mut r, get_weight).unwrap();
627            chosen[item.1] += 1;
628        }
629        verify(chosen);
630
631        // choose_weighted_mut
632        let mut items = [(0u32, 0i32); 14]; // (weight, count)
633        for (i, item) in items.iter_mut().enumerate() {
634            *item = (weights[i], 0);
635        }
636        for _ in 0..N_REPS {
637            items.choose_weighted_mut(&mut r, get_weight).unwrap().1 += 1;
638        }
639        for (ch, item) in chosen.iter_mut().zip(items.iter()) {
640            *ch = item.1;
641        }
642        verify(chosen);
643
644        // Check error cases
645        let empty_slice = &mut [10][0..0];
646        assert_eq!(
647            empty_slice.choose_weighted(&mut r, |_| 1),
648            Err(WeightError::InvalidInput)
649        );
650        assert_eq!(
651            empty_slice.choose_weighted_mut(&mut r, |_| 1),
652            Err(WeightError::InvalidInput)
653        );
654        assert_eq!(
655            ['x'].choose_weighted_mut(&mut r, |_| 0),
656            Err(WeightError::InsufficientNonZero)
657        );
658        assert_eq!(
659            [0, -1].choose_weighted_mut(&mut r, |x| *x),
660            Err(WeightError::InvalidWeight)
661        );
662        assert_eq!(
663            [-1, 0].choose_weighted_mut(&mut r, |x| *x),
664            Err(WeightError::InvalidWeight)
665        );
666    }
667
668    #[test]
669    #[cfg(feature = "std")]
670    fn test_multiple_weighted_edge_cases() {
671        use super::*;
672
673        let mut rng = crate::test::rng(413);
674
675        // Case 1: One of the weights is 0
676        let choices = [('a', 2), ('b', 1), ('c', 0)];
677        for _ in 0..100 {
678            let result = choices
679                .choose_multiple_weighted(&mut rng, 2, |item| item.1)
680                .unwrap()
681                .collect::<Vec<_>>();
682
683            assert_eq!(result.len(), 2);
684            assert!(!result.iter().any(|val| val.0 == 'c'));
685        }
686
687        // Case 2: All of the weights are 0
688        let choices = [('a', 0), ('b', 0), ('c', 0)];
689        let r = choices.choose_multiple_weighted(&mut rng, 2, |item| item.1);
690        assert_eq!(r.unwrap_err(), WeightError::InsufficientNonZero);
691
692        // Case 3: Negative weights
693        let choices = [('a', -1), ('b', 1), ('c', 1)];
694        let r = choices.choose_multiple_weighted(&mut rng, 2, |item| item.1);
695        assert_eq!(r.unwrap_err(), WeightError::InvalidWeight);
696
697        // Case 4: Empty list
698        let choices = [];
699        let r = choices.choose_multiple_weighted(&mut rng, 0, |_: &()| 0);
700        assert_eq!(r.unwrap().count(), 0);
701
702        // Case 5: NaN weights
703        let choices = [('a', f64::NAN), ('b', 1.0), ('c', 1.0)];
704        let r = choices.choose_multiple_weighted(&mut rng, 2, |item| item.1);
705        assert_eq!(r.unwrap_err(), WeightError::InvalidWeight);
706
707        // Case 6: +infinity weights
708        let choices = [('a', f64::INFINITY), ('b', 1.0), ('c', 1.0)];
709        for _ in 0..100 {
710            let result = choices
711                .choose_multiple_weighted(&mut rng, 2, |item| item.1)
712                .unwrap()
713                .collect::<Vec<_>>();
714            assert_eq!(result.len(), 2);
715            assert!(result.iter().any(|val| val.0 == 'a'));
716        }
717
718        // Case 7: -infinity weights
719        let choices = [('a', f64::NEG_INFINITY), ('b', 1.0), ('c', 1.0)];
720        let r = choices.choose_multiple_weighted(&mut rng, 2, |item| item.1);
721        assert_eq!(r.unwrap_err(), WeightError::InvalidWeight);
722
723        // Case 8: -0 weights
724        let choices = [('a', -0.0), ('b', 1.0), ('c', 1.0)];
725        let r = choices.choose_multiple_weighted(&mut rng, 2, |item| item.1);
726        assert!(r.is_ok());
727    }
728
729    #[test]
730    #[cfg(feature = "std")]
731    fn test_multiple_weighted_distributions() {
732        use super::*;
733
734        // The theoretical probabilities of the different outcomes are:
735        // AB: 0.5   * 0.667 = 0.3333
736        // AC: 0.5   * 0.333 = 0.1667
737        // BA: 0.333 * 0.75  = 0.25
738        // BC: 0.333 * 0.25  = 0.0833
739        // CA: 0.167 * 0.6   = 0.1
740        // CB: 0.167 * 0.4   = 0.0667
741        let choices = [('a', 3), ('b', 2), ('c', 1)];
742        let mut rng = crate::test::rng(414);
743
744        let mut results = [0i32; 3];
745        let expected_results = [5833, 2667, 1500];
746        for _ in 0..10000 {
747            let result = choices
748                .choose_multiple_weighted(&mut rng, 2, |item| item.1)
749                .unwrap()
750                .collect::<Vec<_>>();
751
752            assert_eq!(result.len(), 2);
753
754            match (result[0].0, result[1].0) {
755                ('a', 'b') | ('b', 'a') => {
756                    results[0] += 1;
757                }
758                ('a', 'c') | ('c', 'a') => {
759                    results[1] += 1;
760                }
761                ('b', 'c') | ('c', 'b') => {
762                    results[2] += 1;
763                }
764                (_, _) => panic!("unexpected result"),
765            }
766        }
767
768        let mut diffs = results
769            .iter()
770            .zip(&expected_results)
771            .map(|(a, b)| (a - b).abs());
772        assert!(!diffs.any(|deviation| deviation > 100));
773    }
774}