如何使用numpy保留数组中固定数量随机位置的非零值，其余替换为掩码

首先明确需求：

• 处理带分类值的numpy二维数组，0为背景值不做修改
• 所有非0分类各自随机保留指定数量的原值，其余同分类元素替换为0
• 保留位置支持随机生成，也可通过随机种子固定重复结果

2021年9月30日更新：已调整实现代码，循环中新增elif分支处理元素数量小于等于目标保留数的类别，可供后续参考，当前版本部分代码存在冗余可根据需求优化

实现代码 #

import numpy as np

# 自定义每个非0类别需要保留的元素数量
KEEP_CNT = 3

new_array = test_array.copy()
uniques, counts = np.unique(new_array, return_counts=True)
print("classes:", uniques, "counts:", counts)

for unique, count in zip(uniques, counts):
    # 跳过背景类0
    if unique == 0:
        continue
    # 该类元素数量大于保留数时，随机选超出数量的元素置0
    if count > KEEP_CNT:
        ids = np.random.choice(count, count - KEEP_CNT, replace=False)
        new_array[tuple(i[ids] for i in np.where(new_array == unique))] = 0
    # 该类元素数量小于等于保留数时，全部保留原值
    else:
        ids = np.random.choice(count, count, replace=False)
        new_array[tuple(i[ids] for i in np.where(new_array == unique))] = unique

如果需要固定每次运行的结果，在代码开头添加随机种子设置即可，例如：

np.random.seed(42)

测试示例 #

测试输入数组：

test_array = np.array([[0,0,0,0,0],
                      [1,1,1,1,1],
                      [0,0,0,0,0],
                      [2,2,2,4,4],
                      [4,4,4,2,2],
                      [0,0,0,0,0]])
print("existing classes:", np.unique(test_array))
# 输出：existing classes: [0 1 2 4]

将KEEP_CNT设置为2时，两次运行的输出示例：
第一次运行结果：

[[0 0 0 0 0]
 [1 0 0 1 0]
 [0 0 0 0 0]
 [2 0 0 0 4]
 [4 0 0 2 0]
 [0 0 0 0 0]]

第二次运行结果：

[[0 0 0 0 0]
 [1 1 0 0 0]
 [0 0 0 0 0]
 [2 0 2 0 4]
 [4 0 0 0 0]
 [0 0 0 0 0]]

内容的提问来源于stack exchange，提问作者Claire Wang