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



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

ID1, z, e, PA, n

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

ID2, RA, DEC

我想匹配ID1ID2以找到共同的以形成另一个包含ID, z, e, PA, n, RA, DEC的"表"。ID1中的大多数元素都可以在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

好吧,您可以将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]])

假设你的第二个表(表 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 的长度并导致索引错误。

使用熊猫:

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[' z'] 引用第二列。 通过修改 read 语句,不再需要这样做:

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

相关内容

  • 没有找到相关文章

最新更新