从 2D numpy 数组中提取任意小计

Question

我有二维计数数组，需要从中提取任意小计序列。在此示例中，它们是小计列。每个小计是任意基列集合的总和，由加数索引元组表示：

>>> A

[[11, 12, 13, 14, 15]

 [21, 22, 23, 24, 25]

 [31, 32, 33, 34, 35]]

>>> subtotal_addend_idxs

((0, 1), (1, 2, 3), (3, 4))

>>> desired_result

[[23, 39, 29]

 [43, 69, 49]

 [63, 99, 69]]

到目前为止我拥有的最好的代码是这样的：

subtotal_addend_idxs = ((0, 1), (1, 2, 3), (3, 4))

np.hstack(

    tuple(

        np.sum(A[:, subtotal_addend_idxs], axis=1, keepdims=True)

        for addend_idxs in self._column_addend_idxs

    )

)

有没有一种聪明的方法可以通过单个numpy调用/表达式来完成此操作，而无需for循环创建各个小计列的元组？

请注意，加数索引是任意的；并非所有索引都需要出现在小计中，索引不一定按升序出现，并且同一索引可以出现在多个小计中。

Answer 1

尝试np.add.reduceat：lens = [len(n) for n in subtotal_addend_idxs] c = np.concatenate(subtotal_addend_idxs) output = np.add.reduceat(A[:,c], np.cumsum([0]+lens)[:-1], axis=1)输出：array([[23, 39, 29],        [43, 69, 49],        [63, 99, 69]], dtype=int32)备注：更快的选项np.concatenate是np.fromiter(itertools.chain(*subtotal_addend_idxs), dtype=int)。

Answer 2

由于我们不能使用np.take，这是我的解决方案（函数中仍然有一个 for 循环lambda...）test = np.array([[11, 12, 13, 14, 15],                 [21, 22, 23, 24, 25],                 [31, 32, 33, 34, 35]])inds = ((0, 1), (1, 2, 3), (3, 4))fake_take = lambda array,inds:[np.sum(array[list(ind)]) for ind in inds]np.apply_along_axis(lambda x:fake_take(x,inds),1,test)