CODEWITHSUNDEEP
pythonnumpymatrix
Arthur Grigorev
Arthur Grigorev

Reputation: 113

Apply a function to each row of a numpy matrix w.r.t. its index

I have a numpy matrix A of shape [n,m] and an array b of length n. What I need is to take sum of b[i] least elements of the i'th row of A. So the code might look like this:

A = np.array([[1,2,3],
              [4,5,6],
              [7,8,9]])

b = np.array([2,3,1])

sums = magic_function() #sums = [3, 15, 7]

I've considered np.apply_along_axis() function but it seems that your function can only depend on the row itself in this case.

Upvotes: 2

Views: 398

Answers (1)

Divakar
Divakar

Reputation: 221764

Vectorized approach making use of NumPy broadcasting to create the mask of valid ones along each row and then perform sum-reduction -

mask = b[:,None] > np.arange(A.shape[1])
out = (A*mask).sum(1)

Alternatively, with np.einsum to get the reduction -

out = np.einsum('ij,ij->i',A,mask)

We can also use np.matmul/@ notation on Python 3.x -

out = (A[:,None] @ mask[...,None]).squeeze()

Upvotes: 5

Related Questions