基于索引前缀在列中设置值的最有效方法

时间:2016-10-19 14:34:34

标签: python pandas optimization

我有一个这样的数据框:

df = pd.DataFrame(index=['pre1_xyz', 'pre1_foo', 'pre3_bar', 'pre3_foo', 'pre10_foo', 'pre10_bar', 'pre10_xyz'])

我想要添加一个列values,其中使用函数return_something(pref)根据相应行的索引的前缀确定该值。现在我按如下方式实现:

import pandas as pd
import numpy as np

# this just returns a random value for the sake of simplicity
def return_something(pref):

    return np.random.choice(len(pref)+10)


df = pd.DataFrame(index=['pre1_xyz', 'pre1_foo', 'pre3_bar', 'pre3_foo', 'pre10_foo', 'pre10_bar', 'pre10_xyz'])

# get all the unique prefixes
unique_pref = set([pi.partition('_')[0] for pi in df.index])

# determine the value for each prefix
val_pref = {pref: return_something(pref) for pref in unique_pref}

# add the values to the dataframe
for prefi, vali in val_pref.items():

    # determine all rows with the same prefix
    rows = [rowi for rowi in df.index if rowi.startswith(prefi+'_')]

    df.loc[rows, 'values'] = vali

然后给了我想要的输出:

           values
pre1_xyz        0
pre1_foo        0
pre3_bar        7
pre3_foo        7
pre10_foo      13
pre10_bar      13
pre10_xyz      13

问题是,是否有比这更聪明的事情,例如避免创建unique_pref和/或val_pref和/或使用set_value的解决方案,这似乎是向数据框添加值的最快解决方案,如this question中所述

1 个答案:

答案 0 :(得分:3)

由于您重复了前缀,因此您需要先分离出前缀,以确保不为同一前缀生成新的随机数。因此,必须从前缀列表中删除重复项。我通过为前缀创建一个新列然后使用df.prefix.unique()以更简洁的方式完成此操作。

df['prefix'] = [i.split('_')[0] for i in df.index]
df['values'] = df.prefix.map(dict(zip(df.prefix.unique(),[return_something(i) for i in df.prefix.unique()])))