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

Question

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.

Is A sorted by column?

kabanus
– kabanus

2017-11-05 13:05:56 +00:00
Commented Nov 5, 2017 at 13:05 — kabanus
– kabanus, Commented Nov 5, 2017 at 13:05
In general case no.

Arthur Grigorev
– Arthur Grigorev

2017-11-05 13:34:59 +00:00
Commented Nov 5, 2017 at 13:34 — Arthur Grigorev
– Arthur Grigorev, Commented Nov 5, 2017 at 13:34

Divakar · Accepted Answer · 2017-11-05 13:18:04Z

5

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

edited Nov 5, 2017 at 13:18

answered Nov 5, 2017 at 13:07

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

cs95 Over a year ago

I like how there's always an einsum solution in your answers.

Collectives™ on Stack Overflow

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

1 Answer 1

1 Comment

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Related