pandas字符串操作:根据其他列内容动态选择列

时间:2019-08-19 16:30:10

标签: python string pandas dataframe

我有一个具有以下结构的DataFrame:

df = pd.DataFrame({
    'variable': ['foo 1 bar', 'foo 2 bar', 'foo 3 bar'],
    'lookup 1': [1, 2, 3],
    'lookup 2': [11, 12, 13],
    'lookup 3': [21, 22, 23],
})

我想添加一个新列,为每一行选择与lookup <i>相对应的foo <i> bar列:

   lookup 1  lookup 2  lookup 3   variable  result
0         1        11        21  foo 1 bar       1
1         2        12        22  foo 2 bar      12
2         3        13        23  foo 3 bar      23

我一直在尝试通过pandas.Series.str.extract提取i,而该pandas.Series.str.cat仅返回df.lookup1

# doesn't work
df.result = df.loc[:, lambda df: "lookup " + df.variable.str.extract("(\d)")]

尝试用here连接字符串也不能这样:

# doesn't work either
df.result = df.loc[:, lambda x: pd.Series(['lookup '] * x.shape[0]).str.cat(x.variable.str.extract("(\d)"))]

3 个答案:

答案 0 :(得分:2)

还可以使用numpy索引,这可能会更快

r, c = df.shape
mapper = dict(zip(df.columns.str.get(-1), 
                  np.arange(c)))

df.values[np.arange(r), 
          df.variable.str.get(4).map(mapper)]

答案 1 :(得分:1)

尝试一下:

df['result'] = df.lookup(df.index, 
                         ('lookup '+ df.variable.str.extract('(\d+)'))[0].tolist())

输出:

    variable  lookup 1  lookup 2  lookup 3  result
0  foo 1 bar         1        11        21       1
1  foo 2 bar         2        12        22      12
2  foo 3 bar         3        13        23      23

答案 2 :(得分:0)

概括/扩展@rafaelc的解决方案,我最终得到以下结果:

# explicitly define which variable value maps to which column
mapper_name = {'foo 1 bar': 'lookup 1',
               'foo 2 bar': 'lookup 2',
               'foo 3 bar': 'lookup 3'}

# map all names to their corresponding column-index
mapper_col_idx = dict(zip(df.columns, np.arange(df.shape[1])))

# apply both mappings to get a Series of column indices, i.e.
# 'foo 1 bar' --maps-to--> 'lookup 1' --maps-to--> 0
col_idx = df.variable.map(mapper_name).map(mapper_col_idx)

# access the dataframe's array via the col_idx array
df['result'] = df.values[np.arange(df.shape[0]), col_idx]