无须计算就可求和

时间:2019-06-20 13:02:15

标签: python python-3.x numpy

考虑以下两个玩具阵列:

import numpy as np
k = np.random.randint(1, 25, (5, 2, 3))
l = np.random.randint(25, 50, (7, 3))

In [27]: k
Out[27]: 
array([[[14, 15, 24],
        [21, 24,  5]],

       [[22, 19,  9],
        [21,  1, 11]],

       [[ 1, 23,  5],
        [16, 14,  2]],

       [[ 7,  3, 16],
        [23,  2,  8]],

       [[12, 24,  4],
        [ 2, 15, 20]]])

In [28]: l
Out[28]: 
array([[47, 31, 42],
       [28, 27, 26],
       [45, 32, 49],
       [29, 34, 32],
       [40, 36, 25],
       [44, 27, 31],
       [27, 35, 26]])

我可以得到我感兴趣的乘法和,如下所示:

f = np.array([np.sum( k * x, axis = 2) for x in l])

In [29]: f
Out[29]: 
array([[[2131, 1941],
        [2001, 1480],
        [ 970, 1270],
        [1094, 1479],
        [1476, 1399]],

       [[1421, 1366],
        [1363,  901],
        [ 779,  878],
        [ 693,  906],
        [1088,  981]],

       [[2286, 1958],
        [2039, 1516],
        [1026, 1266],
        [1195, 1491],
        [1504, 1550]],

       [[1684, 1585],
        [1572,  995],
        [ 971, 1004],
        [ 817,  991],
        [1292, 1208]],

       [[1700, 1829],
        [1789, 1151],
        [ 993, 1194],
        [ 788, 1192],
        [1444, 1120]],

       [[1765, 1727],
        [1760, 1292],
        [ 820, 1144],
        [ 885, 1314],
        [1300, 1113]],

       [[1527, 1537],
        [1493,  888],
        [ 962,  974],
        [ 710,  899],
        [1268, 1099]]])

如何在不求助于理解的情况下计算该总和?

2 个答案:

答案 0 :(得分:3)

这是np.einsum的好用例:

np.einsum('ijk,lk->lij', k, l)

list_comp = np.array([np.sum( k * x, axis = 2) for x in l])
np.allclose(np.einsum('ijk,lk->lij', k, l), list_comp)
# True

或使用broadcasting

(l[:,None,None]*k).sum(-1)

尽管快速检查了时间np.einsum的速度大约快了3倍

答案 1 :(得分:1)

您也可以使用np.tensordot来做到这一点:

import numpy as np

np.random.seed(0)
k = np.random.randint(1, 25, (5, 2, 3))
l = np.random.randint(25, 50, (7, 3))
f = np.tensordot(l, k, [-1, -1])
f_comp = np.array([np.sum(k * x, axis=2) for x in l])
print(np.allclose(f, f_comp))
# True