Skip to content

Commit 8f17ca6

Browse files
committed
Implement SliceExt::partial_shuffle
1 parent 5400221 commit 8f17ca6

File tree

1 file changed

+33
-1
lines changed

1 file changed

+33
-1
lines changed

src/seq.rs

+33-1
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,23 @@ impl<T> SliceExt for [T] {
271271
fn partial_shuffle<R>(&mut self, rng: &mut R, amount: usize)
272272
-> (&mut [Self::Item], &mut [Self::Item]) where R: Rng + ?Sized
273273
{
274-
unimplemented!()
274+
// This applies Durstenfeld's algorithm for the
275+
// [Fisher–Yates shuffle](https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm)
276+
// for an unbiased permutation, but exits early after choosing `amount`
277+
// elements.
278+
279+
let len = self.len();
280+
let mut i = len;
281+
let end = if amount >= len { 0 } else { len - amount };
282+
283+
while i > end {
284+
// invariant: elements with index > i have been locked in place.
285+
i -= 1;
286+
// lock element i in place.
287+
self.swap(i, rng.gen_range(0, i + 1));
288+
}
289+
let r = self.split_at_mut(i);
290+
(r.1, r.0)
275291
}
276292
}
277293

@@ -531,6 +547,22 @@ mod test {
531547
let b: &[_] = &[1, 1, 1];
532548
assert_eq!(x, b);
533549
}
550+
551+
#[test]
552+
fn test_partial_shuffle() {
553+
let mut r = ::test::rng(118);
554+
555+
let mut empty: [u32; 0] = [];
556+
let res = empty.partial_shuffle(&mut r, 10);
557+
assert_eq!((res.0.len(), res.1.len()), (0, 0));
558+
559+
let mut v = [1, 2, 3, 4, 5];
560+
let res = v.partial_shuffle(&mut r, 2);
561+
assert_eq!((res.0.len(), res.1.len()), (2, 3));
562+
assert!(res.0[0] != res.0[1]);
563+
// First elements are only modified if selected, so at least one isn't modified:
564+
assert!(res.1[0] == 1 || res.1[1] == 2 || res.1[2] == 3);
565+
}
534566

535567
#[test]
536568
#[cfg(feature = "alloc")]

0 commit comments

Comments
 (0)