如何在Python矩阵中随机保留每行仅两个1？

原代码问题分析 #

你当前的代码uvv*=np.cumsum(uvv,axis=1)==1是通过累加每行元素，只保留第一个出现的1，所以只能实现每行留一个1的效果，无法满足随机保留两个1的需求。下面提供两种可行的实现思路：

方法一：逐行随机选择（直观易理解） #

通过遍历每行，找到所有1的位置后随机挑选2个，再构造结果矩阵：

import numpy as np

# 原始矩阵
uvv = np.array([[1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1],
                [0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0],
                [1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1],
                [0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0]])

# 初始化全0结果矩阵
result = np.zeros_like(uvv)

for row_idx in range(uvv.shape[0]):
    # 获取当前行中所有1的索引
    one_positions = np.where(uvv[row_idx] == 1)[0]
    # 随机选择2个不重复的索引
    selected_pos = np.random.choice(one_positions, size=2, replace=False)
    # 将选中位置设为1
    result[row_idx, selected_pos] = 1

print(result)

方法二：向量化实现（适合大矩阵，效率更高） #

避免循环，用向量化操作完成随机选择，适合处理大规模数据：

import numpy as np

uvv = np.array([[1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1],
                [0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0],
                [1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1],
                [0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0]])

# 生成仅1的位置有随机权重的矩阵，0的位置权重为0
random_weights = np.random.rand(*uvv.shape) * (uvv == 1)
# 对每行取权重最大的2个位置的索引
top2_indices = np.argpartition(random_weights, -2, axis=1)[:, -2:]

# 构造结果矩阵
result = np.zeros_like(uvv)
rows = np.arange(uvv.shape[0])[:, np.newaxis]
result[rows, top2_indices] = 1

print(result)

内容的提问来源于stack exchange，提问作者Maluma Corazón