问题描述
让a = np.arange(1,4)
。
为了得到 a
的二维乘法表,我这样做:
>>> a * a[:,None]
>>> array([[1,2,3],[2,4,6],[3,6,9]])
对于 3 个维度,我可以执行以下操作:
>>> a * a[:,None] * a[:,None,None]
>>> array([[[ 1,[ 2,[ 3,9]],[[ 2,[ 4,8,12],[ 6,12,18]],[[ 3,9],18],[ 9,18,27]]])
我如何编写一个函数,以一个 numpy 数组 a
和多个维度 n
作为输入并输出 n
的 a
维乘法表?
解决方法
这应该可以满足您的需求:
import itertools
a = np.arange(1,4)
n = 3
def f(x,y):
return np.expand_dims(x,len(x.shape))*y
l = list(itertools.accumulate(np.repeat(np.atleast_2d(a),n,axis=0),f))[-1]
只需将 n
更改为您需要的任何维度
首先,我们可以使用 numpy.expand_dims()
动态提升列表/生成器推导式中所需的数组维度,然后在 Python 3.8 上使用可迭代的产品工具,例如 math.prod
+。实现将如下所示:
from math import prod
def n_dim_multiplication(arr,num_dims):
gen_arr = (np.expand_dims(a,axis=tuple(range(1,idx+1))) for idx in range(num_dims))
return prod(gen_arr)
3 维情况的样本运行:
# input array
In [83]: a = np.arange(1,4)
# desired number of dimensions
In [84]: num_dims = 3
In [85]: n_dim_multiplication(a,num_dims)
Out[85]:
array([[[ 1,2,3],[ 2,4,6],[ 3,6,9]],[[ 2,[ 4,8,12],[ 6,12,18]],[[ 3,9],18],[ 9,18,27]]])