您可以通过广播做到这一点:
import numpy as np
a = np.array([1, 2, 3])
b = np.array([4, 5])
c = a[None, ...] * b[..., None]
print(c)
输出:
[[ 4 8 12]
[ 5 10 15]]
这可以通过制作适当的切片传递给操作数来轻松概括。
编辑
这种泛化的实现可能是:
import numpy as np
def apply_multi_broadcast_1d(func, dim1_arrs):
n = len(dim1_arrs)
iter_dim1_arrs = iter(dim1_arrs)
slicing = tuple(
slice(None) if j == 0 else None
for j in range(n))
result = next(iter_dim1_arrs)[slicing]
for i, dim1_arr in enumerate(iter_dim1_arrs, 1):
slicing = tuple(
slice(None) if j == i else None
for j in range(n))
result = func(result, dim1_arr[slicing])
return result
dim1_arrs = [np.arange(1, n + 1) for n in range(2, 5)]
print(dim1_arrs)
# [array([1, 2]), array([1, 2, 3]), array([1, 2, 3, 4])]
arr = apply_multi_broadcast_1d(lambda x, y: x * y, dim1_arrs)
print(arr.shape)
# (2, 3, 4)
print(arr)
# [[[ 1 2 3 4]
# [ 2 4 6 8]
# [ 3 6 9 12]]
# [[ 2 4 6 8]
# [ 4 8 12 16]
# [ 6 12 18 24]]]
这里不需要递归,我不确定它有什么好处。
另一种方法是从 Python 函数生成np.ufunc(如@TlsChris's answer 中建议的那样)并使用其np.ufunc.outer() 方法:
import numpy as np
def apply_multi_outer(func, dim1_arrs):
ufunc = np.frompyfunc(func, 2, 1)
iter_dim1_arrs = iter(dim1_arrs)
result = next(iter_dim1_arrs)
for dim1_arr in iter_dim1_arrs:
result = ufunc.outer(result, dim1_arr)
return result
虽然这会给出相同的结果(对于一维数组),但它比广播方法要慢(从轻微到显着取决于输入大小)。
此外,虽然apply_multi_broadcast_1d() 仅限于一维输入,但apply_multi_outer() 也适用于更高维的输入数组。广播方法可以很容易地适应更高维度的输入,如下所示。
编辑 2
将apply_multi_broadcast_1d() 推广到 N-dim 输入,包括将广播与函数应用程序分离,如下所示:
import numpy as np
def multi_broadcast(arrs):
for i, arr in enumerate(arrs):
yield arr[tuple(
slice(None) if j == i else None
for j, arr in enumerate(arrs) for d in arr.shape)]
def apply_multi_broadcast(func, arrs):
gen_arrs = multi_broadcast(arrs)
result = next(gen_arrs)
for i, arr in enumerate(gen_arrs, 1):
result = func(result, arr)
return result
三者的基准测试表明apply_multi_broadcast() 比apply_multi_broadcast_1d() 稍慢但比apply_multi_outer() 快:
def f(x, y):
return x * y
dim1_arrs = [np.arange(1, n + 1) for n in range(2, 5)]
print(np.all(apply_multi_outer(f, dim1_arrs) == apply_multi_broadcast_1d(f, dim1_arrs)))
print(np.all(apply_multi_outer(f, dim1_arrs) == apply_multi_broadcast(f, dim1_arrs)))
# True
# True
%timeit apply_multi_broadcast_1d(f, dim1_arrs)
# 100000 loops, best of 3: 7.76 µs per loop
%timeit apply_multi_outer(f, dim1_arrs)
# 100000 loops, best of 3: 9.46 µs per loop
%timeit apply_multi_broadcast(f, dim1_arrs)
# 100000 loops, best of 3: 8.63 µs per loop
dim1_arrs = [np.arange(1, n + 1) for n in range(10, 16)]
print(np.all(apply_multi_outer(f, dim1_arrs) == apply_multi_broadcast_1d(f, dim1_arrs)))
print(np.all(apply_multi_outer(f, dim1_arrs) == apply_multi_broadcast(f, dim1_arrs)))
# True
# True
%timeit apply_multi_broadcast_1d(f, dim1_arrs)
# 100 loops, best of 3: 10 ms per loop
%timeit apply_multi_outer(f, dim1_arrs)
# 1 loop, best of 3: 538 ms per loop
%timeit apply_multi_broadcast(f, dim1_arrs)
# 100 loops, best of 3: 10.1 ms per loop