python - 沿 numpy 数组的一个轴随机保留不为零的单个元素

Question

我需要一种有效的方法来创建形状为 (x,y,3) 的 numpy 数组，其中每个元组 (x,y) 的 3 个随机元素中只有一个随机元素具有从 [-1,0 中随机选择的值,1]

 np.random.randint(-1, 2, (x,y,3))

仅满足我的要求的后半部分。

我可以使用嵌套循环来迭代每个 (x, y) 并将其值乘以随机掩码，但它根本效率不高。

这是循环实现:

a=np.random.randint(-1, 2, (x,y,3))
for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        mask = np.array(np.random.permutation([0,1,0]))
        a[i][j] = a[i][j] * mask

Answer 1

我不会生成一大堆额外的数字并关闭其中的大部分，而是从只生成您需要的数字的角度来解决这个问题。您想要为每个 x-y 对分配一个介于 0 和 2 之间的随机索引。因此生成一个随机索引和随机值，并分配:

indices = np.random.randint(3, size=(x, y))
values = np.random.randint(-1, 2, size=(x, y))
result = np.zeros((x, y, 3), dtype=int)
result[(*np.ogrid[:x, :y], indices)] = values

索引表达式是 advanced index因为 indices 是一个整数列表。在这种情况下，对前两个索引使用 ... 或 :, : 不会达到您想要的效果。相反，np.ogrid生成正确形状的范围，以强制indices 的元素对应于正确的 x-y 坐标。