如何使用numpy有效地从每列中减去值

Question

I have a 2D array of shape (50,50). I need to subtract a value from each column of this array skipping the first), which is calculated based on the index of the column. For example, using a for loop it would look something like this:

for idx in range(1, A[0, :].shape[0]):
    A[0, idx] -= idx * (...) # simple calculations with idx 

Now, of course this works fine, but it's very slow and performance is critical for my application. I've tried computing the values to be subtracted using np.fromfunction() and then subtracting it from the original array, but results are different than those obtained by the for loop iteractive subtraction:

 func = lambda i, j: j * (...) #some simple calculations
 subtraction_matrix = np.fromfunction(np.vectorize(func), (1,50))

 A[0, 1:] -= subtraction_matrix

What am I doing wrong? Or is there some other method that would be better? Any help is appreciated!

Answer 1

All your code snippets indicate that you require the subtraction to happen only in the first row of A (though you've not explicitly mentioned that). So, I'm proceeding with that understanding.

Referring to your use of from_function() , you can use the subtraction_matrix as below:

A[0,1:] -= subtraction_matrix[1:]

Testing it out (assuming shape (5,5) instead of (50,50) ):

import numpy as np

A = np.arange(25).reshape(5,5)
print (A)

func = lambda j: j * 10 #some simple calculations
subtraction_matrix = np.fromfunction(np.vectorize(func), (5,), dtype=A.dtype)

A[0,1:] -= subtraction_matrix[1:]
print (A)

Output:

[[ 0  1  2  3  4]        # print(A), before subtraction
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

[[  0  -9 -18 -27 -36]   # print(A), after subtraction
 [  5   6   7   8   9]
 [ 10  11  12  13  14]
 [ 15  16  17  18  19]
 [ 20  21  22  23  24]]

---

If you want the subtraction to happen in all the rows of A , you just need to use the line A[:,1:] -= subtraction_matrix[1:] , instead of the line A[0,1:] -= subtraction_matrix[1:]