如何有效地使用 numpy 面罩？

Question

I have this array 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]]

What I would like to do is... well let me show you.

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

What the above code does is:

Suppose for each row i , the last element is v_i . For each row i , we have a row with the same 2nd element as i and with 3rd element equal to 1. Call this row j . Now we subtract the last element of j from the last element of i . That is, v_i = v_i - v_j .

The code I have pasted above works fine. But it takes way too long (on my actual array which is way bigger than the one I've pasted as an example). I am quite sure it is the list comprehension that is slowing it down. So I am looking for a way to do this faster, possibly even without a loop (or a mask either).

I would also like to ask if there is a way to get the sum of the last elements of the rows, grouped by the value of the 2nd element.

So, for example the first element of that result would be 159.6652+144.8538+195.6714=500.1904

And I would have 4 such numbers. Again, I have done this using a loop but it takes too much time to run!

I am new to numpy and just learned how important speed is when working with large datasets. I would be grateful if I can learn something new from here. Thanks for taking the time to read this. Please feel free to comment if anything isn't clear.

Answer 1

Here is a solution using np.unique . It makes no assumption on the order of rows. If the 2nd column is already grouped and ordered as in your example this can be simplified.

# 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])