在pandas中按列组合过滤

时间:2016-03-16 14:09:09

标签: python pandas

我有两个DataFrames“A”和“B”。每个都有两列“key1”和“key2”,但唯一键是两个的组合。我想从第二个DataFrame中选择包含在DataFrame“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 s

解决方案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 s

这有效,但有更优雅(快速)的解决方案吗?

1 个答案:

答案 0 :(得分:0)

您可以尝试查看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

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

通常,索引上的操作往往比常规DataFrame列上的操作更快。但不管怎样,你可能在这里寻找边际改善。我认为你不能以微秒级的速度完成这项任务。