有没有办法广播布尔掩码？

Question

我正在尝试根据搜索距离减少计算次数。我有N个节点和一个[NxN]布尔掩码，它告诉我哪些节点在具有T真值的其他节点的X距离内。

我还为每个节点提供了[Nx(d)]个数据，其中(d)可以是(1)，(3)或(3x3)。我想要“稀疏”格式，这是一个[Tx(d)]数组，所以我可以沿0轴进行矢量化计算。现在我这样做：

sparseData=data.repeat(data.shape[0],axis=0).reshape(np.concatenate(([data.shape[0],data.shape])))[mask]

哪个有效，但如果N太大，会导致内存错误，原因是[NxNx(d)]数组我用.repeat创建了吗？有没有办法广播这个？如果我这样做：

data[None,...][mask]

它不起作用，但似乎必须有一种更有效的方法来做到这一点。

Answer 1

您可以使用numpy.broadcast_to：

制作视图，而不是重复数据

sparseData = np.broadcast_to(data, (data.shape[0],) + data.shape)[mask]

但是，更容易根据索引选择数据行：

I, J = np.nonzero(mask)
sparseData = data[I]  # could also use J