Finding elements of a list in an array in Python (Numpy, Pytorch)

I am dealing with the following problem: I have a list, let's say a=[1, 2, 3] and an array b=[[2, 4, 6],[3, 2, 5],[4, 1, 3]], the number of elements in the list is equal to the number of rows in the array, the output should be a list c=[-1, 1, 2], where -1 indicates that the number a[0] was not found in the 0-th row of the array b, the other numbers tells the index where the element a[i] is in the row b[i, :], the index enumeration starts at 0.

I solved it using loops but I am searching for a more optimal way to do it using np.where or the torch alternative, I would appreciate any help.

Try this

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)

Or

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)

Output

[-1, 1, 2]

Assuming at most one match per row, you could use:

out = np.full_like(a, -1)
idx, vals = np.where(b == a[:,None])
out[idx] = vals

Output: array([-1, 1, 2])

Reproducible input:

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