🔒 Set all elements to zero after a certain index

[eternal basin]

Assume I have the following array

arr = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6]])

And an array with indices, that has as many elements as arr has rows , i.e. idx.shape[0] == arr.shape[0]

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

I want all elements, after that index that appears in idxs to be set to zero in arr:

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

Since I have to do this for more than 10k elements, I want to avoid a for-loop

[ashen pagoda]

@eternal basin

Python help channel opened

Remember to:

• Ask your Python question, not if you can ask or if there's an expert who can help.
• Show a code sample as text (rather than a screenshot) and the error message, if you've got one.
• Explain what you expect to happen and what actually happens.

:warning: Do not pip install anything that isn't related to your question, especially if asked to over DMs.

Closes after a period of inactivity, or when you send !close.

[warm marlin]

You can achieve this by iterating over the rows of arr and setting the elements after the index specified in idxs to zero. Here's a Python code that does that:

`
import numpy as np
arr = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6]])
idxs = np.array([2, 1, 0])

for i in range(arr.shape[0]):
arr[i, idxs[i] + 1:] = 0

print(arr)`

[eternal basin]

i cannot use a for loop, that takes too long unfortunately

I forgot to mention this, it must be very efficient

[warm marlin]

You can use vectorized operations in NumPy then if you don't want a for loop

[eternal basin]

yeah, but how would you do that in this case

[native adder]

Why also if u take transpose i think its enough to just replace the last n rows

[eternal basin]

What do you mean with why?

Because otherwise it will take too long 😅

[native adder]

Make a function that does it for 1 arr than add the @vectorize decorator?

[warm marlin]

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

arr[mask] = 0

print(arr)`

[native adder]

Whats too long and why do u have a time limit?

[eternal basin]

There is no time limit per se, but why would I choose a solution that takes long if s shorter one exists?

I have to use this operation many times for experiments, thus I would like to save time whenever possible

[native adder]

Bc the short one is complex and not flexible

[native adder]

U could do it onece and save in pkl

If always on the same data

[eternal basin]

the data changes consistently, so that is not an option, but I get your point

[native adder]

For tests prob doing it on an older is fine

[eternal basin]

right

[unreal escarp]

!e ```py
import numpy as np

arr = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6]])
idxs = np.array([[2], [1], [0]]) # or just reshape a standard array
mask = np.arange(4) > idxs

arr[mask] = 0

print(arr)

[ashen pagoda]

:white_check_mark: Your 3.12 eval job has completed with return code 0.

001 | [[1 2 3 0]
002 |  [2 3 0 0]
003 |  [3 0 0 0]]

[unreal escarp]

I believe this is what you want

the key is to create a mask that looks like this

array([[False, False, False,  True],
       [False, False,  True,  True],
       [False,  True,  True,  True]])
``` by just using the default behavior of `[0,1,2,3] > [[2], [1], [0]]`

[eternal basin]

yup, I did something similar now, thanks! I jsut had the issue that I had to index the columns, not the rows, so it became a little bit ugly

i had to repeeat the np.arange, I am sure there is also a better way 😅

[ashen pagoda]

Python help channel closed

This help channel has been closed and it's no longer possible to send messages here. If your question wasn't answered, feel free to create a new post in #1035199133436354600. To maximize your chances of getting a response, check out this guide on asking good questions.