两个 1D numpy / torch 数组的特殊索引以生成另一个数组

Posted 2023-03-29

【中文标题】两个 1D numpy / torch 数组的特殊索引以生成另一个数组

【英文标题】：Special indexing of two 1D numpy / torch arrays to produce another array

【发布时间】：2021-10-30 09:12:14

【问题描述】：

在给定输入数组 x 和 y 时，在 Numpy 或 PyTorch 中寻找一种有效（矢量化）的方式来实现数组 z：

一维数组x 包含一个递增的 ID 列表，每个 ID 重复 1 次或多次（不必为每个 ID 重复相同的次数）。例如，[0 0 0 1 1 2 2 2 2]

0 和 1 的一维数组 y。对于x 中的每个唯一ID，至少有一个元素等于“1”。例如，[1 1 0 1 1 0 0 1 0]。

一维输出数组z 等于y，但在x 中的每个ID 中只保留y 中第一次出现的“1”。该 ID 的 y 的其余元素应设置为“0”。所以在这个例子中，结果是[1 0 0 1 0 0 0 1 0]

x: [0 0 0 1 1 2 2 2 2]
y: [1 1 0 1 1 0 0 1 0]
z: [1 0 0 1 0 0 0 1 0]

我觉得在 Numpy 或 PyTorch 中有一种快速的方法可以做到这一点，但我想不通。

编辑：这是使用 while 循环的“慢”版本

x = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2])
y = np.array([1, 1, 0, 1, 1, 0, 0, 1, 0])
z = y.copy()
n = z.shape[0]
i = 0
while i < n:
    if y[i] == 1:
        current_id = x[i]
        i += 1
        while i < n and x[i] == current_id:
            z[i] = 0
            i += 1
    else:
        i += 1

【问题讨论】：

让我们从缓慢但简单的方法开始吧！

更新问题以包含基本算法

【参考方案1】：

如果您真的想以矢量化方式执行此操作，您可以这样做：

x = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2, 2])         
y = torch.tensor([1, 1, 0, 1, 1, 0, 0, 1, 0])         
                                                      
unique_vals = x.unique_consecutive().unsqueeze(1)     
masked_x = x.masked_fill(y != 1, unique_vals[-1][0]+1)
indices = (unique_vals == masked_x).int().argmax(1)   
res = torch.zeros_like(y).index_fill(0, indices, 1)   
print(res) 

但我不相信它会比天真的版本更快。

【参考方案2】：

你可以使用np.unique:

unq, ind = np.unique(np.stack((x, y)), axis=1, return_index=True)

ind 现在包含每个唯一元素组合的第一次出现。您只需删除 y 为零的那些：

keep = unq[1, :] != 0
ind = ind[keep]

现在你可以直接制作z了：

z = np.zeros_like(y)
z[ind] = 1