def cartesian_jit(arrays):
"""
Generate a cartesian product of input arrays.
Parameters
----------
arrays : list or tuple of arrays
1-D arrays to form the cartesian product of.
Returns
-------
out : ndarray
2-D array of shape (M, len(arrays)) containing cartesian products
formed of input arrays.
Examples
--------
>>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
array([[1, 4, 6],
[1, 4, 7],
[1, 5, 6],
[1, 5, 7],
[2, 4, 6],
[2, 4, 7],
[2, 5, 6],
[2, 5, 7],
[3, 4, 6],
[3, 4, 7],
[3, 5, 6],
[3, 5, 7]])
"""
n = 1
for x in arrays:
n *= x.size
out = np.zeros((n, len(arrays)))
for i in range(len(arrays)):
m = int(n / arrays[i].size)
out[:n, i] = np.repeat(arrays[i], m)
n //= arrays[i].size
n = arrays[-1].size
for k in range(len(arrays)-2, -1, -1):
n *= arrays[k].size
m = int(n / arrays[k].size)
for j in range(1, arrays[k].size):
out[j*m:(j+1)*m,k+1:] = out[0:m,k+1:]
return out
试试Numpy, 请参考https://gist.github.com/hernamesbarbara/68d073f551565de02ac5gist.github.com