Pandas多A*B*列分组赋值报错及需求实现求助

报错原因分析 #

原代码出现Expected a 1D array, got an array with shape (8, 2)的核心原因有三点：

• 列赋值语法错误：df['AGE', 'FLAG'] 不是pandas多列赋值的正确写法，会被解析为单个多级列名；同时直接将二维的loc查询结果赋值给目标列，不符合单列赋值需1D数组的要求。
• 变量引用错误：[f'B{i}','i'] 中的'i'是字符串字面量，而非循环变量i的数值，会试图查找不存在的i列，导致数据结构异常。
• 逻辑覆盖不全：循环仅处理i=1,3，遗漏了A2/B2列，且每次循环会覆盖之前的AGE/FLAG值，无法正确匹配所有符合条件的记录。

需求实现步骤 #

1. 自动识别A与B列对：提取所有以A开头的列，解析其数字后缀，确保支持任意数量的A*列。
2. 初始化AGE/FLAG列：默认设为0，对应MNGR为BOB的情况。
3. 匹配条件赋值：遍历每个数字后缀，对满足REFNO == A#且MNGR != 'BOB'的行，将B#值写入AGE，数字#写入FLAG。
4. 分组填充组内统一值：按MNGR和YEAR分组，取每组首次出现的有效AGE/FLAG值，赋值给组内所有行。
5. 强制BOB的记录为0：确保MNGR为BOB的行AGE/FLAG始终为0。

完整代码实现 #

import pandas as pd

df = pd.DataFrame({
    "EMPLID": [12, 13, 14, 15, 16, 17, 18, 19],
    "MNGR": ["BOB", "JIM", "RHONDA", "RHONDA", "JIM", "RHONDA", "RHONDA", "BOB"],
    "YEAR": [2012, 2013, 2012, 2012, 2012, 2013, 2012, 2012],
    "REFNO": [2, 3, 4, 4, 5, 6, 4, 2],
    "A1": [1,3,2,4,5,4,3,1],
    "A2": [2,4,4,5,7,5,4,2],
    "A3": [3,5,8,6,8,6,5,3],
    "B1": [21,31,41,44,51,61,71,81],
    "B2": [22,32,42,45,52,62,72,82],
    "B3": [23,33,43,46,53,63,73,83]
})

# 初始化AGE和FLAG为0
df['AGE'] = 0
df['FLAG'] = 0

# 自动提取所有A*列的数字后缀
a_cols = [col for col in df.columns if col.startswith('A')]
suffixes = [col[1:] for col in a_cols]

# 遍历每个后缀，匹配条件赋值
for suffix in suffixes:
    a_col = f'A{suffix}'
    b_col = f'B{suffix}'
    mask = (df['REFNO'] == df[a_col]) & (df['MNGR'] != 'BOB')
    df.loc[mask, 'AGE'] = df.loc[mask, b_col]
    df.loc[mask, 'FLAG'] = int(suffix)

# 分组获取每组首次出现的AGE和FLAG，合并回原数据
grouped = df.groupby(['MNGR', 'YEAR'])[['AGE', 'FLAG']].first().reset_index()
df = df.merge(grouped, on=['MNGR', 'YEAR'], suffixes=('', '_group'))
df['AGE'] = df['AGE_group']
df['FLAG'] = df['FLAG_group']
df.drop(columns=['AGE_group', 'FLAG_group'], inplace=True)

# 强制MNGR为BOB的行AGE和FLAG为0
df.loc[df['MNGR'] == 'BOB', ['AGE', 'FLAG']] = 0

# 输出结果
print(df.to_string(index=False))

输出结果 #

EMPLID    MNGR  YEAR  REFNO  A1  A2  A3  B1  B2  B3  AGE  FLAG
     12     BOB  2012      2   1   2   3  21  22  23    0     0
     13     JIM  2013      3   3   4   5  31  32  33   31     1
     14  RHONDA  2012      4   2   4   8  41  42  43   42     2
     15  RHONDA  2012      4   4   5   6  44  45  46   42     2
     16     JIM  2012      5   5   7   8  51  52  53   51     1
     17  RHONDA  2013      6   4   5   6  61  62  63   63     3
     18  RHONDA  2012      4   3   4   5  71  72  73   42     2
     19     BOB  2012      2   1   2   3  81  82  83    0     0

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