How to select elements in numpy array with given starting point indices

Question

For example, I have a matrix like this:

In [2]: a = np.arange(12).reshape(3, 4) In [3]: a Out[3]: array([[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]])

and a starting point index array:

In [4]: idx = np.array([1, 2, 0]) In [5]: idx Out[5]: array([1, 2, 0])

Are there any vectorized ways to do such things:

for i in range(3): # The following are some usecases a[i, idx[i]:] = 0 a[i, idx[i]-1:] = 0 a[i, :idx[i]] = 0 a[i, idx[i]:idx[i]+2] = 0

Edit: expected output:

array([[ 0, x, x, x], [ 4, 5, x, x], [ x, x, x, x]])

x is placeholder indicating what I'd like to select.

Your code returns a as all zeros. Are you looking for something like this? — Brenlla
– Brenlla, Commented Jul 20, 2018 at 11:25
@Brenlla My requirement is similar to what you provided, but that only applies to square matrix. — dianlujitao
– dianlujitao, Commented Jul 21, 2018 at 3:40

Brenlla · Accepted Answer · 2018-07-21 08:00:21Z

This aproach works for rectangular matrices too. Create a boolean mask trough broadcasting:

a = np.arange(12).reshape(3, 4) idx = np.array([1, 2, 0]) mask=np.arange(a.shape[1]) >= idx[:,None] mask #array([[False, True, True, True], # [False, False, True, True], # [ True, True, True, True]], dtype=bool)

Make your placeholder -1, for example, and set the values of a where mask is true equal to that placeholder:

x = -1 a[mask] = x a #array([[ 0, -1, -1, -1], # [ 4, 5, -1, -1], # [-1, -1, -1, -1]])

Thanks, works as expected, didn't realize arange should be on dim 1

Ersel Er · Accepted Answer · 2018-07-20 11:45:23Z

An expected, exact, output is not provided. So, I think the followings may help you in general.

>>> a = np.arange(12).reshape(3, 4) >>> a array([[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]]) >>> row = np.array([0, 1, 2]) >>> col = np.array([1, 2, 0]) >>> a[row, col] array([1, 6, 8])

You can set the row and cols of a to an value:

>>> a[row, col] = 0 >>> a array([[ 0, 0, 2, 3], [ 4, 5, 0, 7], [ 0, 9, 10, 11]])

Collectives™ on Stack Overflow

How to select elements in numpy array with given starting point indices

2 Answers 2

1 Comment

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Linked

Related