[FIXED] How change nth non-zero element to zero in each row in numpy array

Issue

I have a wide, binary, 2-d numpy array as follows:

np_var:
0, 0, 1, 0, 1, ..., 0, 1
1, 0, 1, 0, 0, ..., 1, 0
...

Each row has 8 non-zero elements. I would like to quickly replace the nth non-zero element in each row with a zero (to end up with 7 non-zero elements per row).

Is there a numpy way to perform this replacement quickly without a loop?

Solution

You can get the indices of non zero elements and use them to replace the values in the array

arr = np.array(...)
print(arr)

# [[1 1 1 0 0 1 0 1 1 0 0 1 1 0]
#  [0 1 1 1 1 0 1 0 0 1 1 0 1 0]
#  [1 0 1 1 0 1 1 1 0 1 0 0 1 0]
#  [0 1 1 0 1 0 0 1 1 1 1 0 1 0]
#  [1 1 1 0 0 1 1 0 0 1 1 0 0 1]
#  [0 0 1 1 1 1 1 0 1 1 0 0 1 0]
#  [1 0 1 0 1 0 1 1 1 0 0 1 0 1]
#  [1 0 1 1 1 0 1 1 0 0 1 0 1 0]
#  [0 0 1 1 1 1 0 1 0 1 1 0 0 1]
#  [0 1 1 1 0 0 0 1 1 0 1 1 1 0]]

nth_element = 5
non_zero_count = int(np.count_nonzero(arr) / len(arr)) # can be replaced by 8 if the size is fixed
indices = arr.nonzero()[1][nth_element - 1::non_zero_count]
arr[np.arange(len(arr)), indices] = 5
print(arr)

# [[1 1 1 0 0 1 0 5 1 0 0 1 1 0]
#  [0 1 1 1 1 0 5 0 0 1 1 0 1 0]
#  [1 0 1 1 0 1 5 1 0 1 0 0 1 0]
#  [0 1 1 0 1 0 0 1 5 1 1 0 1 0]
#  [1 1 1 0 0 1 5 0 0 1 1 0 0 1]
#  [0 0 1 1 1 1 5 0 1 1 0 0 1 0]
#  [1 0 1 0 1 0 1 5 1 0 0 1 0 1]
#  [1 0 1 1 1 0 5 1 0 0 1 0 1 0]
#  [0 0 1 1 1 1 0 5 0 1 1 0 0 1]
#  [0 1 1 1 0 0 0 1 5 0 1 1 1 0]]

Answered By - Guy