检查numpy数组中的每个元素是否在另一个数组中

时间:2013-04-11 02:34:29

标签: python numpy

这个问题似乎很容易,但我找不到一个好看的解决方案。我有两个numpy数组(A和B),我想获得A的索引,其中A的元素在B中,并且得到A的索引,其中元素不在B中。

所以,如果

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

目前我正在使用

C = np.searchsorted(A,B)

利用A按顺序排列的事实,并为[1, 3, 5]提供A中元素的索引。这很棒,但如何获得D = [0,2,4,6]A中不在B的元素索引?

4 个答案:

答案 0 :(得分:35)

如果不是B中的每个元素都在A中,

searchsorted可能会给出错误的答案。您可以使用numpy.in1d

A = np.array([1,2,3,4,5,6,7])
B = np.array([2,4,6,8])
mask = np.in1d(A, B)
print np.where(mask)[0]
print np.where(~mask)[0]

输出是:

[1 3 5]
[0 2 4 6]

但是in1d()使用sort,这对于大型数据集来说很慢。如果数据集很大,可以使用pandas:

import pandas as pd
np.where(pd.Index(pd.unique(B)).get_indexer(A) >= 0)[0]

以下是时间比较:

A = np.random.randint(0, 1000, 10000)
B = np.random.randint(0, 1000, 10000)

%timeit np.where(np.in1d(A, B))[0]
%timeit np.where(pd.Index(pd.unique(B)).get_indexer(A) >= 0)[0]

输出:

100 loops, best of 3: 2.09 ms per loop
1000 loops, best of 3: 594 µs per loop

答案 1 :(得分:6)

import numpy as np

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

D = np.delete(np.arange(np.alen(A)), C)

D
#array([0, 2, 4, 6])

答案 2 :(得分:3)

import numpy as np

a = np.array([1, 2, 3, 4, 5, 6, 7])
b = np.array([2, 4, 6])
c = np.searchsorted(a, b)
d = np.searchsorted(a, np.setdiff1d(a, b))

d
#array([0, 2, 4, 6])

答案 3 :(得分:3)

A的元素也在B中:

设置(A)&组(B)

A中不在B中的元素:

set(A) - set(B)