<p>通过显式指定复合数据类型,可以避免numpy试图查找catchall数据类型时出现的一些问题:</p>
<p>代码+一些时间安排:</p>
<pre><code>import numpy as np
import itertools
def cartesian_product_mixed_type(*arrays):
arrays = *map(np.asanyarray, arrays),
dtype = np.dtype([(f'f{i}', a.dtype) for i, a in enumerate(arrays)])
out = np.empty((*map(len, arrays),), dtype)
idx = slice(None), *itertools.repeat(None, len(arrays) - 1)
for i, a in enumerate(arrays):
out[f'f{i}'] = a[idx[:len(arrays) - i]]
return out.ravel()
a = np.arange(4)
b = np.arange(*map(ord, ('A', 'D')), dtype=np.int32).view('U1')
c = np.arange(2.)
np.set_printoptions(threshold=10)
print(f'a={a}')
print(f'b={b}')
print(f'c={c}')
print('itertools')
print(list(itertools.product(a,b,c)))
print('numpy')
print(cartesian_product_mixed_type(a,b,c))
a = np.arange(100)
b = np.arange(*map(ord, ('A', 'z')), dtype=np.int32).view('U1')
c = np.arange(20.)
import timeit
kwds = dict(globals=globals(), number=1000)
print()
print(f'a={a}')
print(f'b={b}')
print(f'c={c}')
print(f"itertools: {timeit.timeit('list(itertools.product(a,b,c))', **kwds):7.4f} ms")
print(f"numpy: {timeit.timeit('cartesian_product_mixed_type(a,b,c)', **kwds):7.4f} ms")
a = np.arange(1000)
b = np.arange(1000, dtype=np.int32).view('U1')
print()
print(f'a={a}')
print(f'b={b}')
print(f"itertools: {timeit.timeit('list(itertools.product(a,b))', **kwds):7.4f} ms")
print(f"numpy: {timeit.timeit('cartesian_product_mixed_type(a,b)', **kwds):7.4f} ms")
</code></pre>
<p>样本输出:</p>
^{pr2}$