Merge pull request #11056 from bashtage/random-permute-perf

MAINT: Improve performance of random permutation

3 files changed, 37 insertions, 4 deletions

diff --git a/benchmarks/benchmarks/bench_random.py b/benchmarks/benchmarks/bench_random.py
index 7ed3e2fa1..9d84d83d3 100644
--- a/benchmarks/benchmarks/bench_random.py
+++ b/benchmarks/benchmarks/bench_random.py
@@ -65,3 +65,18 @@ class Randint_dtype(Benchmark):
         high = self.high[name]
         np.random.randint(0, high + 1, size=10**5, dtype=name)
 
+
+class Permutation(Benchmark):
+    def setup(self):
+        self.n = 10000
+        self.a_1d = np.random.random_sample(self.n)
+        self.a_2d = np.random.random_sample((self.n, 2))
+    
+    def time_permutation_1d(self):
+        np.random.permutation(self.a_1d)
+
+    def time_permutation_2d(self):
+        np.random.permutation(self.a_2d)        
+
+    def time_permutation_int(self):
+        np.random.permutation(self.n)
diff --git a/doc/release/1.15.0-notes.rst b/doc/release/1.15.0-notes.rst
index 348ca9fc0..23fca0596 100644
--- a/doc/release/1.15.0-notes.rst
+++ b/doc/release/1.15.0-notes.rst
@@ -299,6 +299,10 @@ of the overlap between input an output, that is, the next element accumulated
 is added before the accumulated result is stored in its place, hence the
 overlap is safe. Avoiding the copy results in faster execution.
 
+Increased performance in ``random.permutation`` for multidimensional arrays
+---------------------------------------------------------------------------
+``permutation`` uses the fast path in ``random.shuffle`` for all input
+array dimensions.  Previously the fast path was only used for 1-d arrays.
 
 Changes
 =======
diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx
index 8ef153c15..b45b3146f 100644
--- a/numpy/random/mtrand/mtrand.pyx
+++ b/numpy/random/mtrand/mtrand.pyx
@@ -4901,10 +4901,24 @@ cdef class RandomState:
         """
         if isinstance(x, (int, long, np.integer)):
             arr = np.arange(x)
-        else:
-            arr = np.array(x)
-        self.shuffle(arr)
-        return arr
+            self.shuffle(arr)
+            return arr
+
+        arr = np.asarray(x)
+    
+        # shuffle has fast-path for 1-d
+        if arr.ndim == 1:
+            # must return a copy
+            if arr is x:
+                arr = np.array(arr)
+            self.shuffle(arr)
+            return arr
+
+        # Shuffle index array, dtype to ensure fast path
+        idx = np.arange(arr.shape[0], dtype=np.intp)
+        self.shuffle(idx)
+        return arr[idx]
+        
 
 _rand = RandomState()
 seed = _rand.seed