数组的 numpy 和反对角线

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 中最有效 and/or 最干净的方法是什么？

编辑以添加：根据@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 和反对角线

numpy sum antidiagonals of array

python

arrays

numpy

scipy

multidimensional-array