ENH Speed up and simplify cartesian product

Author: stefanv

sklearn/utils/extmath.py

@@ -6,6 +6,7 @@
 #          Alexandre T. Passos
 #          Olivier Grisel
 #          Lars Buitinck
+#          Stefan van der Walt
 # License: BSD 3 clause
 
 from functools import partial
@@ -501,24 +502,20 @@ def cartesian(arrays, out=None):
            [3, 5, 6],
            [3, 5, 7]])
 
-    References
-    ----------
-    http://stackoverflow.com/q/1208118
-
     """
-    arrays = [np.asarray(x).ravel() for x in arrays]
+    arrays = [np.asarray(x) for x in arrays]
+    shape = (len(x) for x in arrays)
     dtype = arrays[0].dtype
 
-    n = np.prod([x.size for x in arrays])
+    ix = np.indices(shape)
+    ix = ix.reshape(len(arrays), -1).T
+
     if out is None:
-        out = np.empty([n, len(arrays)], dtype=dtype)
-
-    m = n // arrays[0].size
-    out[:, 0] = np.repeat(arrays[0], m)
-    if arrays[1:]:
-        cartesian(arrays[1:], out=out[0:m, 1:])
-        for j in xrange(1, arrays[0].size):
-            out[j * m:(j + 1) * m, 1:] = out[0:m, 1:]
+        out = np.empty_like(ix, dtype=dtype)
+
+    for n, arr in enumerate(arrays):
+        out[:, n] = arrays[n][ix[:, n]]
+
     return out