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

Question

在给定输入数组 x 和 y:

时，在 Numpy 或 PyTorch 中寻找一种有效的（矢量化）方法来实现数组 z

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

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

1D 输出数组 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

Answer 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)

但我不相信它会比原始版本更快。

Answer 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

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

Special indexing of two 1D numpy / torch arrays to produce another array

arrays

numpy

torch

pytorch