python - 通过掩码选择 numpy 数组的元素并保留原始尺寸

Question

您好，我有以下数据

ids = np.concatenate([1.0 * np.ones(shape=(4, 9,)), 
                      2.0 * np.ones(shape=(4, 3,))], axis=1)

logits = np.random.normal(size=(4, 9 + 3, 256))

现在我只想获取 id 为 1.0 的 numpy 数组，并且我想获取大小为 (4,9, 256)

的数组

我尝试了 logits[ids == 1.0, :] 但我得到了 (36, 256)如何在不连接前两个维度的情况下进行切片？

当前尺寸仅为示例尺寸，我正在寻找通用解决方案。

Answer 1

您的问题似乎假设每一行都有相同数量的非零条目；在这种情况下，您通常可以像这样解决问题:

mask = (ids == 1)
num_per_row = mask.sum(1)

# same number of entries per row is required
assert np.all(num_per_row == num_per_row[0])  

result = logits[mask].reshape(logits.shape[0], num_per_row[0], logits.shape[2])

print(result.shape)
# (4, 9, 256)