python - 在 Python 中查找数组中列表的元素(Numpy、Pytorch)

Question

我正在处理以下问题:我有一个列表，比如说 a=[1, 2, 3] 和一个数组 b=[[2, 4, 6] ,[3, 2, 5],[4, 1, 3]]，列表中的元素数量等于数组的行数，输出应该是一个列表c =[-1, 1, 2]，其中-1表示在第0行没有找到数字a[0]数组 b，其他数字告诉索引元素 a[i] 位于行 b[i, :] 中，索引枚举从 0 开始。

我使用循环解决了它，但我正在寻找一种更优化的方法来使用 np.where 或 torch 替代方案，我将不胜感激。

Answer 1

试试这个

import numpy as np

a = [1, 2, 3]
b = np.array([[2, 4, 6], [3, 2, 5], [4, 1, 3]])

# Create a boolean mask by comparing each element of b with the corresponding element of a
mask = b == np.array(a)[:, np.newaxis]

# Find the index of the first occurrence of True along each row of the mask
indices = np.argmax(mask, axis=1)

# Use np.where to conditionally assign the indices where matches were found and -1 where no matches occurred
c = np.where(np.any(mask, axis=1), indices, -1).tolist()

print(c)

或者

import numpy as np

a = [1, 2, 3]
b = np.array([[2, 4, 6], [3, 2, 5], [4, 1, 3]])

c = []
for i in range(len(a)):
    # Find the indices where the i-th element of a matches the elements in the i-th row of b
    indices = np.where(b[i, :] == a[i])[0]
    if len(indices) > 0:
        # If indices are found, append the first index to the list c
        c.append(indices[0])
    else:
        # If no indices are found, append -1 to the list c
        c.append(-1)

print(c)

输出

[-1, 1, 2]