如何获取DataFrame每行4列窗口最大值对应列在另一DataFrame中的值？

现有两个pandas DataFrame（df和df1），需要完成以下操作：

• 在df的每一行中，将每4列划分为一个窗口
• 找到每个窗口内最大值所在的列
• 提取df1中对应行、对应列的值

示例数据 #

import pandas as pd

df = pd.DataFrame()
df['a'] = [1, 2]
df['b'] = [2, 20]
df['c'] = [4, 2]
df['d'] = [6, 2]
df['e'] = [10, 2]
df['f'] = [3, 2]
df['g'] = [5, 22]
df['h'] = [1, 2]
# df输出：
#    a   b  c  d   e  f  g  h
# 0  1   2  4  6  10  3  5  1
# 1  2  20  2  2   2  2 22  2

df1 = pd.DataFrame()
df1['a'] = [1, 2]
df1['b'] = [2, 2]
df1['c'] = [4, 12]
df1['d'] = [1, 2]
df1['e'] = [-5, 2]
df1['f'] = [3, 32]
df1['g'] = [15, -1]
df1['h'] = [1, 2]
# df1输出：
#    a  b   c  d   e   f   g  h
# 0  1  2   4  1  -5   3  15  1
# 1  2  2  12  2   2  32  -1  2

示例说明 #

以第一行为例：

• 窗口(a,b,c,d)的最大值6在d列，对应df1的d列值为1
• 窗口(e,f,g,h)的最大值10在e列，对应df1的e列值为-5

最终预期输出：

a1  b1
0   1  -5
1   2  -1

解决方案代码 #

import pandas as pd

# 构建示例数据（可替换为实际数据）
df = pd.DataFrame({
    'a': [1, 2], 'b': [2, 20], 'c': [4, 2], 'd': [6, 2],
    'e': [10, 2], 'f': [3, 2], 'g': [5, 22], 'h': [1, 2]
})

df1 = pd.DataFrame({
    'a': [1, 2], 'b': [2, 2], 'c': [4, 12], 'd': [1, 2],
    'e': [-5, 2], 'f': [3, 32], 'g': [15, -1], 'h': [1, 2]
})

# 将df的列按每4个一组拆分
window_col_groups = [df.columns[i:i+4] for i in range(0, len(df.columns), 4)]

result_list = []
for idx, cols in enumerate(window_col_groups):
    # 获取当前窗口每行最大值对应的列名
    max_col_names = df[cols].idxmax(axis=1)
    # 从df1中提取对应行和列的值
    extracted_vals = df1.lookup(df1.index, max_col_names)
    # 转为DataFrame加入结果列表，设置对应列名
    result_list.append(pd.DataFrame(extracted_vals, columns=[f'{chr(ord("a")+idx)}1']))

# 合并所有窗口的结果
final_result = pd.concat(result_list, axis=1)
print(final_result)

关键步骤解释 #

1. 拆分列窗口：用列表推导式将df的列按每4个一组分割，得到每个窗口的列名集合
2. 定位最大值列：对每个窗口的子DataFrame调用idxmax(axis=1)，获取每行最大值所在的列名
3. 提取对应值：利用df1.lookup()方法，根据行索引和列名精准提取df1中的对应值
4. 合并结果：将每个窗口的提取结果合并成一个DataFrame，并设置对应的列名

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