python - 根据索引前缀在列中设置值的最有效方法

Question

我有一个像这样的数据框:

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 中所述。 .

Answer 1

由于前缀有重复，因此您需要首先分离出前缀，以确保不会为同一前缀生成新的随机数。因此，有必要从前缀列表中删除重复项。我通过为前缀创建一个新列然后使用 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()])))