diff --git a/src/seq.rs b/src/seq.rs index b4751f25137..e74ea909c37 100644 --- a/src/seq.rs +++ b/src/seq.rs @@ -271,7 +271,23 @@ impl SliceExt for [T] { fn partial_shuffle(&mut self, rng: &mut R, amount: usize) -> (&mut [Self::Item], &mut [Self::Item]) where R: Rng + ?Sized { - unimplemented!() + // This applies Durstenfeld's algorithm for the + // [Fisher–Yates shuffle](https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm) + // for an unbiased permutation, but exits early after choosing `amount` + // elements. + + let len = self.len(); + let mut i = len; + let end = if amount >= len { 0 } else { len - amount }; + + while i > end { + // invariant: elements with index > i have been locked in place. + i -= 1; + // lock element i in place. + self.swap(i, rng.gen_range(0, i + 1)); + } + let r = self.split_at_mut(i); + (r.1, r.0) } } @@ -531,6 +547,22 @@ mod test { let b: &[_] = &[1, 1, 1]; assert_eq!(x, b); } + + #[test] + fn test_partial_shuffle() { + let mut r = ::test::rng(118); + + let mut empty: [u32; 0] = []; + let res = empty.partial_shuffle(&mut r, 10); + assert_eq!((res.0.len(), res.1.len()), (0, 0)); + + let mut v = [1, 2, 3, 4, 5]; + let res = v.partial_shuffle(&mut r, 2); + assert_eq!((res.0.len(), res.1.len()), (2, 3)); + assert!(res.0[0] != res.0[1]); + // First elements are only modified if selected, so at least one isn't modified: + assert!(res.1[0] == 1 || res.1[1] == 2 || res.1[2] == 3); + } #[test] #[cfg(feature = "alloc")]