从 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)