代码之家 › 专栏 › 技术社区 › Jivan

在构建Pandas数据帧时避免循环

vectorization pandas python

Jivan · 技术社区 · 6 年前

我有一个初始的Pandas数据框,有三列,其中一列包含字符串列表。目标是将每一行拆分为尽可能多的元素,就像 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 有100000个元素,计算起来很容易需要一个小时 transformed .

有没有一种矢量化的方法可以在不使用Python循环的情况下获得相同的结果?

1 回复 | 直到 6 年前

jpp 6 年前

本质上,你是重复或链式根据您的列的值。

所以你可以用 np.repeat 和 itertools.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

推荐文章