如何在Python中让numpy矩阵每列仅保留一个1值？

现有一个6列的NumPy数组：

import numpy as np
u = np.array([[1, 0, 0, 1, 0, 0],
              [0, 0, 1, 0, 1, 0],
              [0, 0, 1, 1, 0, 0],
              [0, 0, 0, 1, 0, 0]])

需要实现让矩阵的每一列仅保留一个1：

• 第1、5列已经符合要求（各只有1个1），无需修改
• 第3、4列存在多个1，需要随机删除多余的1，只留一个
• 第2、6列全为0，无需修改

处理后的预期结果示例如下（因随机选择，结果不唯一）：

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

可以通过遍历数组的每一列，针对不同情况处理：

1. 定位每一列中值为1的索引位置
2. 若该列1的数量大于1，随机挑选一个索引保留1，其余位置设为0
3. 1的数量为0或1的列直接跳过

具体实现代码：

import numpy as np
import random

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

# 遍历每一列
for col_idx in range(u.shape[1]):
    # 获取当前列中值为1的行索引
    one_indices = np.where(u[:, col_idx] == 1)[0]
    # 如果1的数量大于1，随机保留一个
    if len(one_indices) > 1:
        keep_idx = random.choice(one_indices)
        # 先把列内所有1置0
        u[:, col_idx] = 0
        # 仅保留选中位置的1
        u[keep_idx, col_idx] = 1

print(u)

代码说明： #

• 用np.where快速定位1的位置，比逐行遍历效率更高
• random.choice保证了保留位置的随机性
• 每列独立处理，不会影响其他列的原有状态

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