我有一个类似于SQL搜索的任务。我有一个"表",其中包含以下由ID1
标识的一维数组(大约 100 万个元素):
ID1, z, e, PA, n
另一个"表",其中包含由ID2
标识的以下一维数组(约150万个元素):
ID2, RA, DEC
我想匹配ID1
和ID2
以找到共同的以形成另一个包含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')