从一个数据帧中找到另一个数据帧中的元素并返回其索引的快速方法

Question

简单地说，我试图将第一个DataFrame的两列中的值与另一个DataFrame中的相同列中的值进行比较。匹配行的索引作为新列存储在firstDataFrame。你知道吗

让我解释一下：我正在研究地理特征（纬度/经度），主要的DataFrame——叫做df——有55M的观测值，看起来有点像这样：

如您所见，只有两行数据看起来合法（索引2和索引4）。你知道吗

第二个DataFrame——称为legit_df——要小得多，拥有我认为合法的所有地理数据：

不必解释原因，主要任务是将df的每个纬度/经度观测值与legit_df的数据进行比较。当匹配成功时，legit_df的索引被复制到df的新列，结果df如下所示：

值-1用于显示没有成功匹配的时间。在上面的例子中，唯一有效的观察结果是索引2和4处的观察结果，它们在legit_df中的索引1和索引2处找到了匹配项。你知道吗

我目前解决这个问题的方法是使用.apply()。是的，它很慢，但是我找不到一种方法来矢量化下面的函数或者使用Cython来加速它：

def getLegitLocationIndex(lat, long):
    idx = legit_df.index[(legit_df['pickup_latitude'] == lat) & (legit_df['pickup_longitude'] == long)].tolist()
    if (not idx):
        return -1
    return idx[0]

df['legit']  = df.apply(lambda row: getLegitLocationIndex(row['pickup_latitude'], row['pickup_longitude']), axis=1)

由于这段代码在有55M个观测值的DataFrame上非常慢，我的问题是：有没有更快的方法来解决这个问题？你知道吗

我正在分享一个Short, Self Contained, Correct (Compilable), Example来帮助你帮助我想出一个更快的选择：

import pandas as pd
import numpy as np

data1 = { 'pickup_latitude'  : [41.366138,   40.190564,  40.769413],
          'pickup_longitude' : [-73.137393, -74.689831, -73.863300]
        }

legit_df = pd.DataFrame(data1)
display(legit_df)

####################################################################################

observations = 10000
lat_numbers = [41.366138,   40.190564,  40.769413, 10, 20, 30, 50, 60, 80, 90, 100]
lon_numbers = [-73.137393, -74.689831, -73.863300, 11, 21, 31, 51, 61, 81, 91, 101]

# Generate 10000 random integers between 0 and 10
random_idx = np.random.randint(low=0, high=len(lat_numbers)-1, size=observations)
lat_data = []
lon_data = []

# Create a Dataframe to store 10000 pairs of geographical coordinates
for i in range(observations):
    lat_data.append(lat_numbers[random_idx[i]])
    lon_data.append(lon_numbers[random_idx[i]])

df = pd.DataFrame({ 'pickup_latitude' : lat_data, 'pickup_longitude': lon_data })
display(df.head())

####################################################################################

def getLegitLocationIndex(lat, long):
    idx = legit_df.index[(legit_df['pickup_latitude'] == lat) & (legit_df['pickup_longitude'] == long)].tolist()
    if (not idx):
        return -1
    return idx[0]


df['legit']  = df.apply(lambda row: getLegitLocationIndex(row['pickup_latitude'], row['pickup_longitude']), axis=1)
display(df.head())

上面的例子创建的df只有10kobservations，在我的机器上运行大约需要7秒。对于100kobservations，运行大约需要67秒。现在想象一下当我要处理5500万排的时候我的痛苦。。。你知道吗

Answer 1

您可以在公共键上使用DataFrame.merge和how='left'。首先重置legit_df的索引。你知道吗

然后fillna与-1：

df.merge(legit_df.reset_index(), on=['pickup_latitude', 'pickup_longitude'], how='left').fillna(-1)

测试性能：

%%timeit
df['legit']  = df.apply(lambda row: getLegitLocationIndex(row['pickup_latitude'], row['pickup_longitude']), axis=1)

每个回路5.81 s±179 ms（平均±标准偏差7次，每个回路1次）

%%timeit
(df.merge(legit_df.reset_index(),on=['pickup_latitude', 'pickup_longitude'], how='left').fillna(-1))

每个回路6.27 ms±254µs（7次运行的平均值±标准偏差，每个100个回路）

Answer 2

我认为使用合并而不是当前的逻辑可以大大加快速度：

full_df = df.merge(legit_df.reset_index(), how="left", on=["pickup_longitude", "pickup_latitude"])

这将重置引用表的索引，使其成为列并在经度上联接

full_df = full_df.rename(index = str, columns={"index":"legit"})
full_df["legit"] = full_df["legit"].fillna(-1).astype(int)

这将重命名为您要查找的列名，并用-1填充join列中的所有缺失

基准：

旧方法： 5.18 s ± 171 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

新方法： 23.2 ms ± 1.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)