Question

鉴于一个numpy ndarray，我想取前两个轴，并用一个新的轴替换它们，这是它们的反对角线的总和。

特别假设我有变量x，y，z，...，并且我的数组的条目代表概率

array[i,j,k,...] = P(x=i, y=j, z=k, ...)

我想获得

new_array[l,k,...] = P(x+y=l, z=k, ...) = sum_i P(x=i, y=l-i, z=k, ...)

即，new_array[l,k,...]是所有array[i,j,k,...]的总和，i+j=l。

在numpy中执行此操作的最有效和/或最干净的方法是什么？

编辑添加：根据@hpaulj的推荐，这是一个明显的迭代解决方案：

array = numpy.arange(30).reshape((2,3,5))
array = array / float(array.sum()) # make it a probability
new_array = numpy.zeros([array.shape[0] + array.shape[1] - 1] + list(array.shape[2:]))
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        new_array[i+j,...] += array[i,j,...]
new_array.sum() # == 1

Answer 1

有trace函数给出对角线的总和。您可以指定偏移量和2个轴（默认值为0和1）。要获得反对角线，您只需要翻转一维。 np.flipud这样做，虽然它只是[::-1,...]索引。

把它们放在一起，

np.array([np.trace(np.flipud(array),offset=k) for k in range(-1,3)])

与您的new_array匹配。

它仍会循环遍历l的可能值（在本例中为4）。 trace本身已编译。

在这个小案例中，它实际上比你的双循环慢（2x3步）。即使我将flipud移出内循环，它仍然较慢。我不知道这对于更大的阵列是如何扩展的。

进一步矢量化这个问题的部分原因是每个对角线都有不同的长度。

In [331]: %%timeit
array1 = array[::-1]
np.array([np.trace(array1,offset=k) for k in range(-1,3)])
   .....: 
10000 loops, best of 3: 87.4 µs per loop

In [332]: %%timeit 
new_array = np.zeros([array.shape[0] + array.shape[1] - 1] + list(array.shape[2:]))                                                       
for i in range(2):
    for j in range(3):
        new_array[i+j] += array[i,j]
   .....: 
10000 loops, best of 3: 43.5 µs per loop

scipy.sparse具有dia格式，用于存储非零对角线的值。它存储填充的值数组以及偏移量。

array([[12,  0,  0,  0],
       [ 8, 13,  0,  0],
       [ 4,  9, 14,  0],
       [ 0,  5, 10, 15],
       [ 0,  1,  6, 11],
       [ 0,  0,  2,  7],
       [ 0,  0,  0,  3]])
array([-3, -2, -1,  0,  1,  2,  3])

虽然这是解决可变对角线长度问题的一种方法，但在这种情况下我不认为它只需要它们的总和。

numpy sum antidiagonals of array

1 个答案: