Given:
a = [[0, 1], [2, 2], [4, 2]]
b = [[0, 1, 2, 3], [2, 2, 3, 4], [4, 2, 3, 3], [2, 3, 3, 3]]
Solution is:
for (i,j) in zip(a[:, 0], a[:, 1]):
print b[np.logical_and( a[:, 0] == i, a[:, 1] == j)]
Result should be
[[0 1 2 3]]
[[2 2 3 4]]
[[4 2 3 3]]
Is there any solution for this problem without use of the for-loop?
You can use NumPy broadcasting for a vectorized solution, like so -
# 2D mask corresponding to all iterations of :
# "np.logical_and( a[:, 0] == i, a[:, 1] == j)"
mask = (a[:,None,0] == a[:,0]) & (a[:,None,1] == a[:,1])
# Use column indices of valid ones for indexing into b for final output
_,C_idx = np.where(mask)
out = b[C_idx]
Sample run -
In [67]: # Modified generic case
...: a = np.array([[0, 1], [3, 2], [3, 2]])
...: b = np.array([[0, 1, 2, 3], [2, 2, 3, 4], [4, 2, 3, 3], [2, 3, 3, 3]])
...:
...: for (i,j) in zip(a[:, 0], a[:, 1]):
...: print b[np.logical_and( a[:, 0] == i, a[:, 1] == j)]
...:
[[0 1 2 3]]
[[2 2 3 4]
[4 2 3 3]]
[[2 2 3 4]
[4 2 3 3]]
In [68]: mask = (a[:,None,0] == a[:,0]) & (a[:,None,1] == a[:,1])
...: _,C_idx = np.where(mask)
...: out = b[C_idx]
...:
In [69]: out
Out[69]:
array([[0, 1, 2, 3],
[2, 2, 3, 4],
[4, 2, 3, 3],
[2, 2, 3, 4],
[4, 2, 3, 3]])
Given:
a = np.array(([0, 1], [2, 2], [4, 2]))
b = np.array(([0, 1, 2, 3], [2, 2, 3, 4], [4, 2, 3, 3], [2, 3, 3, 3]))
Calculate:
temp = np.in1d(b[:,0], a[:,0]) * np.in1d(b[:,1], a[:,1])
result = b[temp]
print 'result:', result
Output:
result: [[0 1 2 3]
[2 2 3 4]
[4 2 3 3]]
