python - 按 pandas 中的列组合过滤

Question

我有两个数据帧“A”和“B”。每个都有两列“key1”和“key2”，但唯一的键是两者的组合。我想从第二个数据帧中选择数据帧“A”中包含的“key1”和“key2”列组合的所有行。

简单示例:

    A = pd.DataFrame({'a':list(range(20000))*100,
                      'b':np.repeat(list(range(100)),20000)})

    B = pd.DataFrame({'a':list(range(40000))*100,
                      'b':np.repeat(list(range(100)),40000),
                      'c':np.random.randint(4000000, size = 4000000)})

解决方案 1:

    %%time
    A['marker'] = True
    C = B.merge(A, on=['a','b'], how='inner').drop('marker', axis=1)

1.26 秒

解决方案 2:

    %%time
    A['marker'] = A['a'].astype(str) + '_' + A['b'].astype(str)
    B['marker'] = B['a'].astype(str) + '_' + B['b'].astype(str)
    C = B[B.marker.isin(A.marker)]

20.4 秒

这可行，但是有更优雅(且快速)的解决方案吗？

Answer 1

您可以尝试查看pd.MultiIndex并使用多级索引而不是普通/无意义的整数索引。不确定在实际数据中是否会快很多，但稍微修改一下示例数据:

index1 = pd.MultiIndex.from_arrays([range(20000)*100, np.repeat(range(100),20000)]) #former A
index2 = pd.MultiIndex.from_arrays([range(40000)*100, np.repeat(range(100),40000)]) #index of B[['a', 'b']]
s = pd.Series(np.random.randint(4000000, size = 4000000), index=index2) #former B['c']

In [93]: %timeit c = s[index1]
1 loops, best of 3: 803 ms per loop

使用与其原始索引(index2)不同的索引(index1)对s进行索引大致相当于您的合并操作.

通常，对索引的操作往往比对常规 DataFrame 列执行的操作更快。但无论哪种方式，您可能都在寻求边际改进。我认为您无法在微秒级内完成此任务。