randomly change numpy array values

Question

I have an array data.

import numpy as np

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

print(data)

If some non_zero elements (1,2,3,4) have more than 5 counts, I want to randomly extract and keep 5 positions and replace all other positions into 0s.

uniques, counts = np.unique(data, return_counts=True)
for unique, count in zip(uniques, counts):
    print (unique, count)

    if count > 5:       
       
       ids = np.random.choice(range(count), 5, replace=False)

How can I do it?

Answer 1

uniques, counts = np.unique(data, return_counts=True)
for unique, count in zip(uniques, counts):
    print (unique, count)
    if unique != 0 and count > 5:
       ids = np.random.choice(count, count-5, replace=False)
       data[tuple(i[ids] for i in np.where(data == unique))] = 0