有效地收集数据框中的链接

时间:2016-09-28 08:43:40

标签: python pandas

假设我有一个类型为

的数据框
individual,  location, food
         1          A     a
         1          A     b
         1          B     a
         1          A     c
         2          C     a
         2          C     b

个人在位置和食物之间建立联系。我想收集个人的所有链接。也就是说,如果在AB位置观察到某个人,并且abc有(最终)食物,我想链接所有这些位置和食物类型相互对立:

    location   food
           A      a
           A      b
           A      c
           B      a
           B      b
           B      c
           C      a
           C      b

一个 - 非常低效 - 这样做的方式是

import itertools
def foo(group):
    list1 = group.location.unique()
    list2 = group.food.unique()
    return pd.DataFrame(data=list(itertools.product(list1, list2)), columns=['location', 'food'])
df.groupby(df.individual).apply(foo)

有没有更好的方法来完成这项工作?

1 个答案:

答案 0 :(得分:2)

你可以使用numpy的meshgrid来提高效率。

import itertools
import numpy as np
def foo(group):
    list1 = group.location.unique()
    list2 = group.food.unique()
    return pd.DataFrame(data=list(itertools.product(list1, list2)), columns=['location', 'food'])

def bar(group):
    list1 = group.location.unique()
    list2 = group.food.unique()
    product = np.meshgrid(list1, list2)
    # reversing the order is necessary to get the same output as foo
    list3 = np.dstack([product[1], product[0]]).reshape(-1, 2)
    return pd.DataFrame(data=list3, columns=['location', 'food'])

在我的机器上有一个小的,(~20%)加速

In [66]: %timeit df.groupby(df.individual).apply(foo)
100 loops, best of 3: 2.57 ms per loop

In [67]: %timeit df.groupby(df.individual).apply(bar)
100 loops, best of 3: 2.16 ms per loop