如何基于另一DataFrame的条件按组选择pd.DataFrame列并精简代码？

我有两个DataFrame（x和y），其中DF x包含group变量以及val_1、val_2两个数值变量。我需要根据DF y中的条件，为x中的每个组选择val_1或val_2列，最终生成一个包含group列和单一数值列val的新DF z。我已经写出了实现代码，但作为Python新手，想知道能不能让代码更精简。

原实现代码：

import pandas as pd
import numpy as np

# 定义分组和每组观测数
groups = ['a','b','c']
N_group = 5
L = len(groups)*N_group

# 构建DF x 
x = pd.DataFrame({'group': np.repeat(groups, N_group),
                  'val_1': np.random.uniform(size=L),
                  'val_2': np.random.uniform(low=10, high=20,size=L)})

# 条件DF：指定每个分组选择val_1还是val_2
y = pd.DataFrame({'group': groups,
                  'val_col' : ['val_1','val_2','val_1']})

# 用循环根据y的行内容筛选x，能否优化？
z = []
for i, r in y.iterrows():
    z.append(pd.DataFrame({'group': r[0],
                           'val': x.loc[x['group'] == r[0], r[1]]
                           }))
z = pd.concat(z)
print(z)

输出结果：

group        val
0      a   0.346833
1      a   0.575308
2      a   0.083667
3      a   0.878797
4      a   0.649935
5      b  17.904204
6      b  15.361618
7      b  15.837250
8      b  11.352091
9      b  11.390637
10     c   0.123949
11     c   0.022118
12     c   0.737024
13     c   0.230206
14     c   0.613382

精简优化方案 #

方法1：合并后用lookup（最简洁高效） #

先将x和y按group合并，再用pd.DataFrame.lookup直接根据val_col的值提取对应列的数值：

import pandas as pd
import numpy as np

groups = ['a','b','c']
N_group = 5
L = len(groups)*N_group

x = pd.DataFrame({'group': np.repeat(groups, N_group),
                  'val_1': np.random.uniform(size=L),
                  'val_2': np.random.uniform(low=10, high=20,size=L)})
y = pd.DataFrame({'group': groups, 'val_col' : ['val_1','val_2','val_1']})

# 合并两个DataFrame
merged = x.merge(y, on='group')
# 用lookup提取对应列的值并生成目标DF
z = merged.assign(val=merged.lookup(merged.index, merged['val_col']))[['group', 'val']]
print(z)

方法2：map+apply实现映射提取 #

先把y转成分组到列名的字典映射，再对x的每一行根据group匹配对应列提取数值：

# 构建分组到目标列的映射字典
col_map = y.set_index('group')['val_col'].to_dict()
# 提取对应列的值并生成目标DF
z = x.assign(val=x.apply(lambda row: row[col_map[row['group']]], axis=1))[['group', 'val']]

方法3：numpy索引（性能最优，适合大数据量） #

利用numpy的向量化索引避免循环，大幅提升处理效率：

# 构建分组到目标列的映射字典
col_map = y.set_index('group')['val_col'].to_dict()
# 将目标列名转换成数值列的索引（val_1对应0，val_2对应1）
col_indices = x['group'].map(col_map).map({'val_1':0, 'val_2':1})
# 用numpy索引批量提取数值
vals = x[['val_1', 'val_2']].to_numpy()[np.arange(len(x)), col_indices]
# 生成目标DF
z = pd.DataFrame({'group': x['group'], 'val': vals})

以上方法均避免了原代码中的循环与多次concat操作，代码更简洁的同时，性能也有明显提升，其中lookup和numpy索引的方式尤其适合处理大规模数据集。

内容的提问来源于stack exchange，提问作者moooh