随机更改 numpy 数组值

Question

我有一个数组数据。

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)

如果某些非零元素（1,2,3,4）的计数超过5个，我想随机提取并保留5个位置，并将所有其他位置替换为0。

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)

我该怎么做？

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