如何用Numpy生成指定轴子数组至少含1个1的随机0-1数组？

可以通过先随机生成基础数组，再批量修正全0子数组的方式实现，全程用Numpy向量化操作，避免低效循环：

import numpy as np

size = (3, 5, 5)
proba_0 = 0.7
target_axis = 2  # 指定需要保证至少一个1的轴

# 1. 生成初始随机0-1数组
n_positions = np.random.choice([0, 1], size=size, p=[proba_0, 1-proba_0])

# 2. 定位指定轴上的全0子数组
# 沿目标轴求和，和为0则表示该子数组全为0
axis_sum = np.sum(n_positions, axis=target_axis, keepdims=True)
all_zero_mask = axis_sum == 0

# 3. 批量修正全0子数组：给每个全0子数组随机设一个1
if np.any(all_zero_mask):
    # 生成目标轴维度内的随机索引，对应每个全0子数组要设1的位置
    rand_idx = np.random.randint(0, size[target_axis], size=size[:target_axis] + size[target_axis+1:])
    
    # 构造多维索引，定位到所有需要修正的位置
    grid = np.ogrid[tuple(map(slice, size[:target_axis] + size[target_axis+1:]))]
    grid.insert(target_axis, rand_idx)
    
    # 批量赋值为1
    n_positions[tuple(grid)] = 1

关键逻辑说明 #

• 初始数组生成和原方法一致，保留原有的概率分布
• 通过np.sum+掩码快速定位所有不符合要求的子数组，keepdims=True保证掩码形状可与原数组广播
• 用np.ogrid构造多维索引，一次性对所有全0子数组的随机位置赋值，完全规避Python循环的性能损耗

验证：可以通过np.all(np.sum(n_positions, axis=target_axis) > 0)检查是否所有指定轴的子数组都至少包含一个1。

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