如何有效地使用 numpy 掩码？

Question

我有这个数组a:

[[  1.       1.       0.      42.533   43.53   159.6652]
 [  1.       1.       0.      57.122   28.21   144.8538]
 [  1.       1.       1.      86.586   32.37   195.6714]
 [  1.       2.       1.      33.768    4.89    58.5222]
 [  1.       2.       0.      90.336   30.19   195.9074]
 [  1.       2.       0.      57.099   27.16   142.4066]
 [  2.       3.       0.      48.371   19.14   103.0763]
 [  2.       3.       1.      30.82     4.74    50.02  ]
 [  2.       3.       0.      27.147   50.98   142.3491]
 [  2.       4.       0.      27.275   43.79   127.4165]
 [  2.       4.       0.      79.439    8.79   121.7297]
 [  2.       4.       1.      21.747   44.44   121.5951]]

我想做的是……好吧，让我告诉你。

mask = np.array([np.where((a[:, 1]==i[1]) & (a[:, 2]==1)) for i in a])
a[:, -1] -= a[mask][:, 0, 0, -1]

以上代码的作用是：

假设每一行i，最后一个元素是v_i。 对于每一行 i，我们有一行与 i 具有相同的第二个元素，并且第三个元素等于 1。将此行称为 j。 现在我们从 i 的最后一个元素中减去 j 的最后一个元素。也就是说，v_i = v_i - v_j。

我在上面粘贴的代码工作正常。但这需要的时间太长（在我的实际数组上，比我粘贴的示例大得多）。我很确定是列表理解减慢了它的速度。所以我正在寻找一种方法来更快地做到这一点，甚至可能没有循环（或掩码）。

我还想问是否有办法获取行的最后一个元素的总和，按第二个元素的值分组。

因此，例如该结果的第一个元素将是 159.6652+144.8538+195.6714=500.1904

我会有 4 个这样的数字。同样，我使用循环完成了此操作，但运行时间太长！

我是 numpy 新手，刚刚了解到处理大型数据集时速度的重要性。如果我能从这里学到新东西，我将不胜感激！感谢您抽时间阅读。如果有什么不明白的，请随时发表评论。

Answer 1

这是使用np.unique 的解决方案。它不对行的顺序做任何假设。如果第二列已经按照您的示例进行了分组和排序，则可以简化。

# find unique id's and
# idx such that unq[idx] would recover a[:,1]
unq,idx = np.unique(a[:,1],return_inverse=True)
unq
# array([1., 2., 3., 4.])
idx
# array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3])

# find indices of reference rows     
ridx = a[:,2].nonzero()[0]
ridx
# array([ 2,  3,  7, 11])

# extract reference rows (last col only) in order of unq
ref = np.empty(unq.size,a.dtype)
ref[idx[ridx]] = a[ridx,-1]
ref
# array([195.6714,  58.5222,  50.02  , 121.5951])

# subtract reference
# (replace "-" with "-=" to subtract in-place) 
a[:,-1] - ref[idx]
# array([-3.600620e+01, -5.081760e+01,  0.000000e+00,  0.000000e+00,
#         1.373852e+02,  8.388440e+01,  5.305630e+01,  0.000000e+00,
#         9.232910e+01,  5.821400e+00,  1.346000e-01,  0.000000e+00])

# group sums 
np.bincount(idx,a[:,-1])
# array([500.1904, 396.8362, 295.4454, 370.7413])