如何使用python查找数组中的元素列表

Question

I have two sets of data which both have values which refer to part of a larger set of data (points in an unstructured mesh).

The two smaller sets of data contain vectors which have the global id which references the point in the larger set of data. Something like:

Large set of data:

0 0 0
0 0 1
0 1 0
1 0 0
1 1 0
1 0 1
0 1 1
1 1 1 

Smaller sets of data:

A

0 1
3 5
4 5 
6 7 
7 2 

B

0 10
4 12
7 60

The first column in the smaller sets of data is a reference to the line number in the larger set of data. The second column in the smaller set of data are just example data.

It is also worth mentioning that the first column of B is always a subset of the first column of A.

What I need is the row indices of A where the point ids match those in B.

In this case this would be:

ind = [0,2,4]

ie A[ind,0] = B[:,0]

I have managed to do this previously using a loop, but now the datasets are increasing in size to over 10 million and the loop is far too slow. Can anyone suggest any faster methods?

Answer 1

Putting the first column data of B into a set should speed things up. Assuming that A and B are lists of tuples (or lists), try this:

>>> A
[('0', '1'), ('3', '5'), ('4', '5'), ('6', '7'), ('7', '2')]
>>> B
[('0', '10'), ('4', '12'), ('7', '60')]
>>> bkeys=set([i[0] for i in B])
>>> [i for i,v in enumerate(A) if v[0] in bkeys]
[0, 2, 4]