Numpy：查找两个不同数组中值的索引条件（来自R）

Question

我有一个由 3D ndarray X 表示的体积，其值介于 0 和 255 之间，我还有另一个 3D ndarray Y，它是第一个数组的任意掩码，具有值0 或 1。

我想找到 50 个体素的随机样本的指标，该样本在 X 中大于零，即“图像”，在 Y 中等于 1，即“掩码”。

我的经验是使用 R，其中以下方法可行：

idx <- sample(which(X>0 & Y==1), 50)

也许 R 的优势在于我可以线性索引 3D 数组，因为仅在 numpy 中使用单个索引就可以给我一个 2D 矩阵。

我猜它可能涉及numpy.random.choice，但似乎我不能有条件地使用它，更不用说以两个不同的数组为条件了。我应该使用另一种方法吗？

Answer 1

这是一种方法-

N = 50 # number of samples needed (50 for your actual case)

# Get mask based on conditionals
mask = (X>0) & (Y==1)

# Get corresponding linear indices (easier to random sample in next step)
idx = np.flatnonzero(mask)

# Get random sample
rand_idx = np.random.choice(idx, N)

# Format into three columnar output (each col for each dim/axis)
out = np.c_[np.unravel_index(rand_idx, X.shape)]

如果您需要不带替换的随机样本，请使用np.random.choice() 和可选参数replace=False。

示例运行 -

In [34]: np.random.seed(0)
    ...: X = np.random.randint(0,4,(2,3,4))
    ...: Y = np.random.randint(0,2,(2,3,4))

In [35]: N = 5 # number of samples needed (50 for your actual case)
    ...: mask = (X>0) & (Y==1)
    ...: idx = np.flatnonzero(mask)
    ...: rand_idx = np.random.choice(idx, N)
    ...: out = np.c_[np.unravel_index(rand_idx, X.shape)]

In [37]: mask
Out[37]: 
array([[[False,  True,  True, False],
        [ True, False,  True, False],
        [ True, False,  True,  True]],

       [[False,  True,  True, False],
        [False, False, False,  True],
        [ True,  True,  True,  True]]], dtype=bool)

In [38]: out
Out[38]: 
array([[1, 0, 1],
       [0, 0, 1],
       [0, 0, 2],
       [1, 1, 3],
       [1, 1, 3]])

将输出out 与mask 中True 值的位置相关联，以便快速验证。

---

如果您不想扁平化以获得线性索引并直接获取每个暗淡/轴的索引，我们可以这样做 -

i0,i1,i2 = np.where(mask)
rand_idx = np.random.choice(len(i0), N)
out = np.c_[i0,i1,i2][rand_idx]

为了性能，先索引，然后在最后一步与np.c_连接-

out = np.c_[i0[rand_idx], i1[rand_idx], i2[rand_idx]]