如何在Numpy数组的axis=2维度随机选取值为1的索引？

我有一个如下的NumPy数组：

import numpy as np
arr = np.random.randint(2, size=(32, 4, 19))

需要针对axis=2维度，为数组中每个(32,4)位置随机选取一个值为1的索引，最终返回形状为(32,4,1)的数组。

例如原数组前几行如下：

array([[[0, 1, 0, ..., 0, 1, 1],
        [0, 0, 1, ..., 1, 1, 0],
        [0, 0, 0, ..., 1, 0, 0],
        [0, 0, 1, ..., 0, 0, 0]],
        ...
       [[0, 1, 1, ..., 1, 1, 1],
        [1, 0, 1, ..., 0, 1, 0],
        [1, 1, 0, ..., 1, 0, 0],
        [0, 0, 0, ..., 0, 1, 1]]])

期望得到类似如下的结果：

array([[[1],[17],[16],[5]], 
       [[3], ...
        ....
       [[7],[4],[7],[11]]]) 

其中arr[0,0,1] == 1、arr[0,1,17] == 1，以此类推。

方法一：直观循环实现 #

逻辑简单易懂，适合小规模数组：

import numpy as np

# 生成示例数组
arr = np.random.randint(2, size=(32, 4, 19))

# 初始化结果数组
result = np.zeros((32, 4, 1), dtype=int)

for i in range(arr.shape[0]):
    for j in range(arr.shape[1]):
        # 获取当前(i,j)位置所有值为1的索引
        ones_indices = np.argwhere(arr[i,j] == 1).flatten()
        # 随机选择一个索引，处理全0的情况时默认选0
        selected_idx = np.random.choice(ones_indices) if len(ones_indices) > 0 else 0
        result[i,j,0] = selected_idx

方法二：向量化高效实现 #

利用NumPy的向量化特性，避免显式循环，性能更优：

import numpy as np

arr = np.random.randint(2, size=(32, 4, 19))
depth = arr.shape[2]

# 生成与原数组形状一致的索引矩阵
indices = np.tile(np.arange(depth), arr.shape[:2] + (1,))

# 将值为0的位置替换为一个超出范围的数，方便后续过滤
masked_indices = np.where(arr == 1, indices, depth)

# 对每个(i,j)位置的索引进行随机打乱
shuffled = np.apply_along_axis(np.random.permutation, 2, masked_indices)

# 找到每个位置第一个不是超出范围的数，即为随机选中的1的索引
result = np.argmin(shuffled, axis=2, keepdims=True)

# 处理全0的情况（此时argmin会返回depth，这里替换为0）
result[result == depth] = 0

验证结果 #

可以通过以下代码验证结果是否符合要求：

# 随机取几个位置检查
i, j = 5, 2
selected_idx = result[i,j,0]
print(f"arr[{i},{j},{selected_idx}] = {arr[i,j,selected_idx]}")  # 应输出1（除非该位置全0）

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