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

Question

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?

Answer 1

You can find where you don't have zero then create a replacement array base all non_zero + one_zero and do replace like below: (I write small example with three non zero and replace the third nonzero with zero)

row = 5
non_zero = 3
# creating sample array
arr = np.concatenate((np.zeros((row,2)), np.ones((row,non_zero))), axis=1)
print(np.count_nonzero(arr))
#15

# creating replace array
rep = np.array(([1]*(non_zero-1)+[0])*row)
# suffle array
[np.random.shuffle(x) for x in arr]
print(arr)
# [[0. 0. 1. 1. 1.]
#               ^^ thrid nonzero
#  [1. 0. 1. 0. 1.]
#               ^^ thrid nonzero
#  [0. 1. 0. 1. 1.]
#               ^^ thrid nonzero
#  [1. 0. 1. 1. 0.]
#            ^^ thrid nonzero
#  [0. 1. 1. 1. 0.]]
#            ^^ thrid nonzero

arr[np.where(arr!=0)] = rep
print(np.count_nonzero(arr))
# 10

print(arr)
# [[0. 0. 1. 1. 0.]
#               ^^ thrid nonzero to zero
#  [1. 0. 1. 0. 0.]
#               ^^ thrid nonzero to zero
#  [0. 1. 0. 1. 0.]
#               ^^ thrid nonzero to zero
#  [1. 0. 1. 0. 0.]
#           ^^ thrid nonzero to zero
#  [0. 1. 1. 0. 0.]]
#           ^^ thrid nonzero to zero

Answer 2

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]]