Pandas：通过转换修复损坏对角矩阵遇瓶颈

我正在处理一个左侧存在损坏（NaN值）的对角矩阵（需从右往左读取），需要按照以下规则修复矩阵：

• 规则1：找到每行损坏前的第一个有效数值，例如dim1=e的行中，损坏前的数值是H1列的6.1
• 规则2：在无损坏的行（dim1=b、c、d，对应原数据的第0、1、2行）的对应列中，找到与该数值最接近的两个值
• 规则3：取这两行的H0列平均值作为H0_miss，对应列的平均值作为H_max_miss，将这两个新列添加到原数据中

初始代码 #

import pandas as pd
import numpy as np

df_flat = pd.DataFrame({
    'dim1': ['b', 'b', 'b', 'b', 'c', 'c', 'c', 'c', 'd', 'd', 'd', 'd', 'e', 'e', 'e', 'e', 'f', 'f', 'f', 'f', 'g', 'g', 'g', 'g'],
    'dim2': ['H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3'],
    'val': [2, 4, 3, 3.1, 3.2, 3.3, 6, 2.1, 2.2, 4.1, 3.4, 3.5, np.nan, 6.1, 2.3, 1.2, np.nan, np.nan, 2.4, 2.5, np.nan, np.nan, np.nan, 2.5]})

new_df = df_flat.set_index(['dim1', 'dim2']).unstack()
new_df.columns = new_df.columns.droplevel(0)

预期输出 #

df_flat = pd.DataFrame({
'dim1': ['b', 'b', 'b', 'b', 'b', 'b', 'c', 'c', 'c', 'c', 'c', 'c', 'd', 'd', 'd', 'd', 'd', 'd', 'e', 'e', 'e', 'e', 'e', 'e', 'f', 'f', 'f', 'f', 'f', 'f','g', 'g', 'g', 'g', 'g', 'g'],
'dim2': ['H0', 'H1', 'H2', 'H3', 'H0_miss', 'H_max_miss', 'H0', 'H1', 'H2', 'H3', 'H0_miss', 'H_max_miss', 'H0', 'H1', 'H2', 'H3', 'H0_miss', 'H_max_miss', 'H0', 'H1', 'H2', 'H3', 'H0_miss', 'H_max_miss', 'H0', 'H1', 'H2', 'H3', 'H0_miss', 'H_max_miss', 'H0', 'H1', 'H2', 'H3', 'H0_miss', 'H_max_miss'],
'val': [2, 4, 3, 3.1, np.nan, np.nan, 3.2, 3.3, 6, 2.1, np.nan, np.nan, 2.2, 4.1, 3.4, 3.5, np.nan, np.nan, np.nan, 6.1, 2.3, 1.2, 4.05, 2.1, np.nan, np.nan, 2.4, 2.5, 3.2, 2.1, np.nan, np.nan, np.nan, 2.5, 2.6, 2.6]})

new_df = df_flat.set_index(['dim1', 'dim2']).unstack()
new_df.columns = new_df.columns.droplevel(0)

当前卡住的代码片段 #

a = new_df.to_numpy()
m = np.isnan(a)
col = m.argmin(axis=1)

完整解决方案代码 #

基于已有的代码片段，以下是完成修复的完整实现：

import pandas as pd
import numpy as np

# 初始数据加载
df_flat = pd.DataFrame({
    'dim1': ['b', 'b', 'b', 'b', 'c', 'c', 'c', 'c', 'd', 'd', 'd', 'd', 'e', 'e', 'e', 'e', 'f', 'f', 'f', 'f', 'g', 'g', 'g', 'g'],
    'dim2': ['H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3', 'H0', 'H1', 'H2', 'H3'],
    'val': [2, 4, 3, 3.1, 3.2, 3.3, 6, 2.1, 2.2, 4.1, 3.4, 3.5, np.nan, 6.1, 2.3, 1.2, np.nan, np.nan, 2.4, 2.5, np.nan, np.nan, np.nan, 2.5]})

new_df = df_flat.set_index(['dim1', 'dim2']).unstack()
new_df.columns = new_df.columns.droplevel(0)

# 1. 获取每行损坏前的列索引（第一个非NaN列的索引）
a = new_df.to_numpy()
m = np.isnan(a)
col_indices = m.argmin(axis=1)

# 2. 提取无损坏行（所有列都没有NaN的行）
good_rows = new_df[~m.any(axis=1)]

# 3. 为每个损坏行计算H0_miss和H_max_miss
h0_miss_vals = []
hmax_miss_vals = []

for idx in range(len(new_df)):
    row = new_df.iloc[idx]
    # 无损坏行直接添加NaN
    if not m[idx].any():
        h0_miss_vals.append(np.nan)
        hmax_miss_vals.append(np.nan)
        continue
    
    # 获取损坏前列的索引和对应值
    target_col_idx = col_indices[idx]
    target_col = new_df.columns[target_col_idx]
    target_val = row[target_col]
    
    # 找到无损坏行中对应列最接近的两个值
    diffs = np.abs(good_rows[target_col] - target_val)
    top2_closest = good_rows.iloc[diffs.nsmallest(2).index]
    
    # 计算平均值
    h0_avg = top2_closest['H0'].mean()
    hmax_avg = top2_closest[target_col].mean()
    
    h0_miss_vals.append(h0_avg)
    hmax_miss_vals.append(hmax_avg)

# 4. 添加新列到宽表
new_df['H0_miss'] = h0_miss_vals
new_df['H_max_miss'] = hmax_miss_vals

# 5. 转换回扁平格式，匹配预期输出
result_df_flat = new_df.stack().reset_index(name='val')

# 查看结果
print(result_df_flat)

代码解释 #

• 步骤1：通过np.isnan标记NaN值，argmin(axis=1)定位每行第一个非NaN列的索引，即损坏前的列位置。
• 步骤2：筛选出所有列无NaN的行，作为修复的参考基准。
• 步骤3：遍历每行，对损坏行提取目标值，在参考行中找到最接近的两个样本，计算对应列的平均值作为新列值。
• 步骤4-5：将带新列的宽表转换为扁平格式，与预期输出结构一致。

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