匹配两个numpy数组以查找相同的元素

时间:2016-12-01 05:55:02

标签: python arrays numpy compare match

我的任务类似于SQL搜索。我有一张"表"其中包含由ID1标识的以下1D数组(约1百万个元素):

ID1, z, e, PA, n

另一个"表"其中包含由ID2标识的以下1D数组(约150万个元素):

ID2, RA, DEC

我希望匹配ID1ID2以找到常见的形式来形成另一个"表"其中包含ID, z, e, PA, n, RA, DECID1中的大多数元素都可以在ID2中找到,但不是全部,否则我可以使用numpy.in1d(ID1,ID2)来完成它。任何人都有快速的方法来完成这项任务?

例如:

ID1, z, e, PA, n
101, 1.0, 1.2, 1.5, 1.8
104, 1.5, 1.8, 2.2, 3.1
105, 1.4, 2.0, 3.3, 2.8

ID2, RA, DEC
101, 4.5, 10.5
107, 90.1, 55.5
102, 30.5, 3.3
103, 60.1, 40.6
104, 10.8, 5.6

输出应为

ID, z, e, PA, n, RA, DEC
101, 1.0, 1.2, 1.5, 1.8, 4.5, 10.5
104, 1.5, 1.8, 2.2, 3.1, 10.8, 5.6

3 个答案:

答案 0 :(得分:1)

您可以将np.in1d与交换位置一起用于两个数组/表的第一列,这样我们就可以使用两个掩码来索引数组以供选择。然后,只需叠加结果 -

mask1 = np.in1d(a[:,0], b[:,0])
mask2 = np.in1d(b[:,0], a[:,0])
out = np.column_stack(( a[mask1], b[mask2,1:] ))

示例运行 -

In [44]: a
Out[44]: 
array([[ 101. ,    1. ,    1.2,    1.5,    1.8],
       [ 104. ,    1.5,    1.8,    2.2,    3.1],
       [ 105. ,    1.4,    2. ,    3.3,    2.8]])

In [45]: b
Out[45]: 
array([[ 101. ,    4.5,   10.5],
       [ 102. ,   30.5,    3.3],
       [ 103. ,   60.1,   40.6],
       [ 104. ,   10.8,    5.6],
       [ 107. ,   90.1,   55.5]])

In [46]: mask1 = np.in1d(a[:,0], b[:,0])

In [47]: mask2 = np.in1d(b[:,0], a[:,0])

In [48]: np.column_stack(( a[mask1], b[mask2,1:] ))
Out[48]: 
array([[ 101. ,    1. ,    1.2,    1.5,    1.8,    4.5,   10.5],
       [ 104. ,    1.5,    1.8,    2.2,    3.1,   10.8,    5.6]])

答案 1 :(得分:1)

假设您的第二个表(表B)已排序,您可以执行排序查找,然后检查是否实际找到了索引元素:

idx = np.searchsorted(B[:-1, 0], A[:, 0])
found = A[:, 0] == B[idx, 0]
np.hstack((A[found, :], B[idx[found], 1:]))

结果:

array([[ 101. ,    1. ,    1.2,    1.5,    1.8,    4.5,   10.5],
       [ 104. ,    1.5,    1.8,    2.2,    3.1,   10.8,    5.6]])

排除B索引的最后一个元素,以简化A中的项超出B中的最终元素的情况。没有它,返回的索引可能会大于B的长度并导致索引错误。

答案 2 :(得分:1)

使用pandas:

import pandas as pd

id1 = pd.read_csv('id1.txt')
id2 = pd.read_csv('id2.txt')
df = id1.merge(id2.sort_values(by='ID2').drop_duplicates('ID2').rename(columns={'ID2':'ID1'}))
print(df)

产地:

   ID1    z    e   PA    n    RA   DEC
0  101  1.0  1.2  1.5  1.8   4.5  10.5
1  104  1.5  1.8  2.2  3.1  10.8   5.6

对于大型数据集,您可能需要执行以下操作:

# [Optional] sort locations and drop duplicates
id2.sort_values(by='ID2', inplace=True)
id2.drop_duplicates('ID2', inplace=True)

# columns that you are merging must have the same name
id2.rename(columns={'ID2':'ID1'}, inplace=True)

# perform the merge
df = id1.merge(id2)

如果没有drop_duplicates,每个项目都会得到一行:

df = id1.merge(id2.rename(columns={'ID2':'ID1'}))
print(id2)
print(df)

,并提供:

   ID2    RA   DEC
0  101   4.5  10.5
1  107  90.1  55.5
2  102  30.5   3.3
3  103  60.1  40.6
4  104  10.8   5.6
5  103  60.1  40.6
6  104  10.9   5.6
   ID1    z    e   PA    n    RA   DEC
0  101  1.0  1.2  1.5  1.8   4.5  10.5
1  104  1.5  1.8  2.2  3.1  10.8   5.6
2  104  1.5  1.8  2.2  3.1  10.9   5.6

请注意,此解决方案会保留列的不同类型:

>>> id1.ID1.dtype
dtype('int64')
>>> id1[' z'].dtype
dtype('float64')

由于标题行中逗号后面有空格,这些空格成为列名的一部分,因此需要使用id1引用第二列[' ž&#39]。通过修改read语句,不再需要这样做:

>>> id1 = pd.read_csv('id1.txt', skipinitialspace=True)
>>> id1.z.dtype
dtype('float64')