numpy where condition on index

Question

I have a numpy 1D array of numbers representing columns e.g: [0,0,2,1] And a matrix e.g:

[[1,1,1],
[1,1,1],
[1,1,1],
[1,1,1]]

Now I a want to change the values in the matrix to 0 where the column index is bigger than the value given in the 1D array:

[[1,0,0],
[1,0,0],
[1,1,1],
[1,1,0]]

How can I achieve this? I think I need a condition based on the index, not on the value

Explanation: The first row in the matrix has indices [0,0 ; 0,1 ; 0,2] where the second index is the column. For indices 0,0 ; 0,1 and 0,2 the value 0 is given. 1 and 2 are bigger than 0. Thus only 0,0 is not changed to zero.

Column index (2) is bigger than all values in column with index 2. Why don't you change all 1s to 0s then? — zabop
– zabop, Commented Nov 6, 2021 at 20:05
@zabop where the column index is bigger, for the first row [1,1,1] we have number 0. In this row we have the numbers with indices 0,0 ; 0,1 and 0,2 The first index does not have column >0, the others do. The 1 and 2 out of 0,1 and 0,2 are bigger than 0. — Jabi
– Jabi, Commented Nov 6, 2021 at 20:07

mozway · Accepted Answer · 2021-11-06 20:33:06Z

3

Assuming a the 2D array and v the 1D vector, you can create a mask of the same size and use numpy.where:

x,y = a.shape
np.where(np.tile(np.arange(y), (x,1)) <= v[:,None], a, 0)

Input:

a = np.array([[1,1,1],
              [1,1,1],
              [1,1,1],
              [1,1,1]])

v = np.array([0,0,2,1])

Output:

array([[1, 0, 0],
       [1, 0, 0],
       [1, 1, 1],
       [1, 1, 0]])

Intermediates:

>>> np.tile(np.arange(y), (x,1))
[[0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]]

>>> np.tile(np.arange(y), (x,1)) <= v[:,None]
[[ True False False]
 [ True False False]
 [ True  True  True]
 [ True  True False]]

edited Nov 6, 2021 at 20:33

answered Nov 6, 2021 at 20:27

mozway

267k13 gold badges55 silver badges106 bronze badges

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

1 Comment

Jabi Over a year ago

Thanks for your answer, this works! I accepted your answer, for the clear details, easy two liner and clear intermediate steps.

Rodalm · Accepted Answer · 2021-11-06 20:27:17Z

Construct a 2D array whose elements are the corresponding column index, and then mask the elements greater than the corresponding value of the 1D array.

Taking advantage of broadcasting, you can do:

>>> arr = np.ones((4,3))
>>> arr 

array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]])

>>> col_thr_idx = np.array([0,0,2,1])
>>> col_thr_idx

array([0, 0, 2, 1])

>>> mask = np.arange(arr.shape[1])[None,:] > col_thr_idx[:,None]
>>> mask

array([[False,  True,  True],
       [False,  True,  True],
       [False, False, False],
       [False, False,  True]])

>>> arr[mask] = 0
>>> arr

array([[1., 0., 0.],
       [1., 0., 0.],
       [1., 1., 1.],
       [1., 1., 0.]])

Collectives™ on Stack Overflow

numpy where condition on index

2 Answers 2

1 Comment

Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Related