'numpy, how to get elements of an array that form part of a record in another array

import numpy as np

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = np.array([1,7])
c = 2

How could I get elements of B where there exists a record in A where the first column is an element of B and the second column equals c?



Solution 1:[1]

I found method only with for-loop

import numpy as np

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = np.array([1,7])
c = 2

for x in B:
    result = A[ (A[:,0] == x) & (A[:,1] == c) ]
    if len(result) > 0:  # `if result:` doesn't work because it is `numpy.array`
        print(x, '=>', result)

Result:

1 => [[1 2 3]]

EDIT:

Without for-loop

import numpy as np

A = np.array([[1,2,3],[4,5,6],[7,3,9]])
B = np.array([1,7])
c = 2

result = B[ A[ (A[:,1] == c), 0 ] == B ]

print(result)

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 furas