exp/rand: add Shuffle, copying CL51891 from math/rand
Other than fixing the examples for the different output from the
different generator, and using Int31n instead of int31n, this is
a straight patch of https://golang.org/cl/51891
to bring the packages closer to alignment.
Change-Id: I0a6ca1e2da8209b4da0a5cf03f83b2a6d750ef64
Reviewed-on: https://go-review.googlesource.com/122541
Reviewed-by: Ian Lance Taylor <iant@golang.org>
diff --git a/rand/example_test.go b/rand/example_test.go
index 3b119fc..47f21fd 100644
--- a/rand/example_test.go
+++ b/rand/example_test.go
@@ -7,6 +7,7 @@
import (
"fmt"
"os"
+ "strings"
"text/tabwriter"
"golang.org/x/exp/rand"
@@ -100,3 +101,34 @@
// Uint64n(10) 2 9 4
// Perm [1 3 4 0 2] [2 4 0 3 1] [3 2 0 4 1]
}
+
+func ExampleShuffle() {
+ words := strings.Fields("ink runs from the corners of my mouth")
+ rand.Shuffle(len(words), func(i, j int) {
+ words[i], words[j] = words[j], words[i]
+ })
+ fmt.Println(words)
+
+ // Output:
+ // [ink corners of from mouth runs the my]
+}
+
+func ExampleShuffle_slicesInUnison() {
+ numbers := []byte("12345")
+ letters := []byte("ABCDE")
+ // Shuffle numbers, swapping corresponding entries in letters at the same time.
+ rand.Shuffle(len(numbers), func(i, j int) {
+ numbers[i], numbers[j] = numbers[j], numbers[i]
+ letters[i], letters[j] = letters[j], letters[i]
+ })
+ for i := range numbers {
+ fmt.Printf("%c: %c\n", letters[i], numbers[i])
+ }
+
+ // Output:
+ // D: 4
+ // A: 1
+ // E: 5
+ // B: 2
+ // C: 3
+}
diff --git a/rand/rand.go b/rand/rand.go
index d4001fd..884e3e4 100644
--- a/rand/rand.go
+++ b/rand/rand.go
@@ -185,6 +185,31 @@
return m
}
+// Shuffle pseudo-randomizes the order of elements.
+// n is the number of elements. Shuffle panics if n < 0.
+// swap swaps the elements with indexes i and j.
+func (r *Rand) Shuffle(n int, swap func(i, j int)) {
+ if n < 0 {
+ panic("invalid argument to Shuffle")
+ }
+
+ // Fisher-Yates shuffle: https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
+ // Shuffle really ought not be called with n that doesn't fit in 32 bits.
+ // Not only will it take a very long time, but with 2³¹! possible permutations,
+ // there's no way that any PRNG can have a big enough internal state to
+ // generate even a minuscule percentage of the possible permutations.
+ // Nevertheless, the right API signature accepts an int n, so handle it as best we can.
+ i := n - 1
+ for ; i > 1<<31-1-1; i-- {
+ j := int(r.Int63n(int64(i + 1)))
+ swap(i, j)
+ }
+ for ; i > 0; i-- {
+ j := int(r.Int31n(int32(i + 1)))
+ swap(i, j)
+ }
+}
+
// Read generates len(p) random bytes and writes them into p. It
// always returns len(p) and a nil error.
// Read should not be called concurrently with any other Rand method.
@@ -270,6 +295,11 @@
// from the default Source.
func Perm(n int) []int { return globalRand.Perm(n) }
+// Shuffle pseudo-randomizes the order of elements using the default Source.
+// n is the number of elements. Shuffle panics if n < 0.
+// swap swaps the elements with indexes i and j.
+func Shuffle(n int, swap func(i, j int)) { globalRand.Shuffle(n, swap) }
+
// Read generates len(p) random bytes from the default Source and
// writes them into p. It always returns len(p) and a nil error.
// Read, unlike the Rand.Read method, is safe for concurrent use.
diff --git a/rand/rand_test.go b/rand/rand_test.go
index c81293f..b310fa6 100644
--- a/rand/rand_test.go
+++ b/rand/rand_test.go
@@ -449,6 +449,14 @@
}
}
+func TestShuffleSmall(t *testing.T) {
+ // Check that Shuffle allows n=0 and n=1, but that swap is never called for them.
+ r := New(NewSource(1))
+ for n := 0; n <= 1; n++ {
+ r.Shuffle(n, func(i, j int) { t.Fatalf("swap called, n=%d i=%d j=%d", n, i, j) })
+ }
+}
+
// Benchmarks
func BenchmarkSource(b *testing.B) {
@@ -520,6 +528,30 @@
}
}
+func BenchmarkPerm30ViaShuffle(b *testing.B) {
+ r := New(NewSource(1))
+ for n := b.N; n > 0; n-- {
+ p := make([]int, 30)
+ for i := range p {
+ p[i] = i
+ }
+ r.Shuffle(30, func(i, j int) { p[i], p[j] = p[j], p[i] })
+ }
+}
+
+// BenchmarkShuffleOverhead uses a minimal swap function
+// to measure just the shuffling overhead.
+func BenchmarkShuffleOverhead(b *testing.B) {
+ r := New(NewSource(1))
+ for n := b.N; n > 0; n-- {
+ r.Shuffle(52, func(i, j int) {
+ if i < 0 || i >= 52 || j < 0 || j >= 52 {
+ b.Fatalf("bad swap(%d, %d)", i, j)
+ }
+ })
+ }
+}
+
func BenchmarkRead3(b *testing.B) {
r := New(NewSource(1))
buf := make([]byte, 3)
