构建Pandas DataFrame时避免循环

时间:2018-07-28 10:32:45

标签: python pandas vectorization

我有一个包含三列的初始Pandas DataFrame,其中一列包含一个字符串列表。目标是将每一行划分为与obj列中的项目一样多的元素,例如:

from    to      obj
--------------------
abc     xyz     [foo, bar]
def     uvw     [gee]
ghi     rst     [foo, bar, baz]

成为这个:

from    to      obj
--------------------
abc     xyz     foo
abc     xyz     bar
def     uvw     gee
ghi     rst     foo
ghi     rst     bar
ghi     rst     baz

目前我正在这样做:

transformed = pd.DataFrame(columns=['from', 'to', 'obj'])

for index, row in origin.iterrows():
    for obj in row['obj']:
        transformed = transformed.append(pd.Series({
            'from':     row['from'],
            'to':       row['to'],
            'obj':      obj
        }), ignore_index=True)

这工作得很好,但速度很慢。如果origin有100,000个元素,则很容易需要一个小时来计算transformed

是否有一种矢量化的方法来获得相同的结果,而不必求助于Python循环?

1 个答案:

答案 0 :(得分:1)

本质上,您是根据列重复链接

因此,您可以根据需要使用np.repeatitertools.chain。如您的示例所示,该解决方案对于少量的列是有效的。

import numpy as np
from itertools import chain

# set up dataframe
df = pd.DataFrame({'from': ['abc', 'def', 'gfhi'],
                   'to': ['xyz', 'uvw', 'rst'],
                   'obj': [['foo', 'bar'], ['gee'], ['foo', 'bar', 'baz']]})

# calculate length of each list in obj
lens = df['obj'].map(len)

# calculate result, repeating or chaining as appropriate
res = pd.DataFrame({'from': np.repeat(df['from'], lens),
                    'to': np.repeat(df['to'], lens),
                    'obj': list(chain.from_iterable(df['obj']))})

print(res)

   from   to  obj
0   abc  xyz  foo
0   abc  xyz  bar
1   def  uvw  gee
2  gfhi  rst  foo
2  gfhi  rst  bar
2  gfhi  rst  baz