'How to return values from an array when they match certain values in a different array

I have two arrays, a and b, where a is a 2 dimensional array and b is a 3 dimensional array. Where the values in a row of array a match the first two values in a row of array b, I want to return the whole row from array b.

a = np.array([[1,1], [3,4],[8,7],[9,10]])    
b = np.array([[0,1,3], [3,4,5],[6,7,8],[9,10,11]])

So for these arrays, I would want to return [3,4,5] and [9,10,11]. I'm not totally sure where to start with this so any advice helps, thank you so much.



Solution 1:[1]

The question isn't completely clear, but you could try:

result = b[np.equal(a, b[:, :2]).all(axis=1)]

Result for your example:

array([[ 3, 4, 5], [ 9, 10, 11]])

Or without Numpy:

result = [
    r2 for r1, r2 in zip(a, b) if all(i1 == i2 for i1, i2 in zip(r1, r2))
]

Result:

[array([3, 4, 5]), array([ 9, 10, 11])]

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 Timus