How to efficiently find the indices a first array values matching with a second array values?

Viewed 134

I have two numpy arrays A and B. A has shape (10000000, 3) and B has shape (1000000, 3). Both the arrays are XYZ coordinates such that B corresponds to some region of A. I have to find indexes of A which correspond to values B. Right now I am solving as below. I would like some help in optimizing this using Numpy or other python packages.

extract_BinA=np.empty(B.shape[0])
for i in range(B.shape[0]):
    for j in range(A.shape[0]):
        if(A[j][0]==B[i][0] and A[j][1]==B[i][1] and A[j][2]==B[i][2]):
            extract_BinA[i]=j
2 Answers

The issue here is not the speed of pure-python code, but the algorithm itself. You can use sorted-arrays or hash-tables to improve the complexity of the algorithm to O(n log n) or even O(n) rather than the slow current O(n^2) solution (as well as the solution proposed by @Mazen). An O(n^2) cannot be efficient here since it will results in roughly 10,000,000 * 10,000,000 = 100,000 billion operations which is too much for any modern computer.

Here is a hash-table solution in pure Python:

table = {tuple(A[i]):i for i in range(A.shape[0])}
extract_BinA = np.empty(B.shape[0])
for i in range(B.shape[0]):
    val = tuple(B[i])
    if val in table:
        extract_BinA[i] = table[val]

Note that the result may differ if there are multiple points in the same location in A.

Here is a benchmark with two random array of size 10,000:

Initial solution: 53.82 s
Mazen solution:    1.76 s
This solution:     0.02 s

On this small input, the above code is 2700 times faster than the initial solution and 88 times faster than the proposed alternative solution. On bigger input, the gap will be much bigger and the above code is many order of magnitude faster than the two other solutions (ie. >10000 times faster).


Update:

If there are multiple points equal each other in A, then the dictionary can be modified to store list of indices rather than one value. Alternatively, the dictionary can be created so that the first value is kept like in the original code. Here are example of the two solutions:

table = dict()
for i in range(A.shape[0])
    key = tuple(A[i])
    if key in table:
        table[key].append(i)
    else:
        table[key] = [i]

extract_BinA = np.empty(B.shape[0])
for i in range(B.shape[0]):
    val = tuple(B[i])
    if val in table:
        # Here table[val] is a list and thus you 
        # can do whatever you want with the indices. 
        # For example you can take the first one like here, 
        # or possibly the last as you want.
        extract_BinA[i] = table[val][0]
# Select always directly the first index
table = dict()
for i in range(A.shape[0])
    key = tuple(A[i])
    if key not in table:
        table[key] = i

extract_BinA = np.empty(B.shape[0])
for i in range(B.shape[0]):
    val = tuple(B[i])
    if val in table:
        extract_BinA[i] = table[val]

Note that these solution are a bit slower than the above code but the complexity is still linear (and thus still very fast).

Solution

extract_BinA = np.ones(B.shape[0]) * -1
for i, b in enumerate(B):
    idx = np.argwhere((A == b) == [True, True, True])
    if idx.any():
        extract_BinA[i] = idx[0][0]
print(extract_BinA)

Explanation

Set extract_BinA to negative values array of size B

extract_BinA = np.ones(B.shape[0]) * -1

In order to get the indices of the elements where B elements equals to A elements, we would need to do the following:

(A == b)

Compares x,y,z for a row in B with every x,y,z rows in A

(A == b) == [True, True, True]

Compares only elements where x_a==x_b, y_a==y_b, and z_a==z_b yields True for all of them

np.argwhere((A == b) == [True, True, True])

Returns a set of indices where the condition is true

A full example to test:

import numpy as np
A = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12],[13,14,15]])
B = np.array([[0,0,0],[4,5,6],[13,14,15]])
# your code
extract_BinA=np.ones(B.shape[0]) * -1
for i in range(B.shape[0]):
    for j in range(A.shape[0]):
        if (A[j] == B[i]).all():
            extract_BinA[i]=j
print(extract_BinA)
# my code
extract_BinA = np.ones(B.shape[0]) * -1
for i, b in enumerate(B):
    idx = np.argwhere((A == b) == [True, True, True])
    if idx.any():
        extract_BinA[i] = idx[0][0] -----------> changed  extract_BinB to  extract_BinA
print(extract_BinA)
Related