如何基于条件填充空DataFrame？需用for循环实现

原始数据 #

现有condition DataFrame如下：

[0]   [1]   [2]   [3]
1     0     0     1     0     
2     0     1     0     0     
3     0     0     0     1     
4     0     0     0     1     

可通过以下代码复现：

import numpy as np
import pandas as pd
n = 4
t = 3   
condition = pd.DataFrame([[0,0,1,0], [0,1,0,0], [0,0,0, 1], [0,0,0, 1]], columns=['0','1', '2', '3'])
condition.index = np.arange(1, n+1)

需求说明 #

需要通过for循环填充df，逻辑如下：

• 逐元素检查condition：
  1. 若condition值为1：df对应位置填充df_3对应值的平方（df_3所有值为3）
  2. 若condition值为0：
     • 该行当前列之后的所有列无1 → 填充2.5
     • 该行当前列之后存在1 → 填充5

期望得到的df结果：

[0]   [1]   [2]   [3]
1     5     5     9    2.5     
2     5     9    2.5   2.5     
3     5     5     5     9     
4     5     5     5     9   

错误代码问题点 #

原代码存在以下问题：

• 循环变量命名冲突（t被重复用作循环变量）
• 条件判断语法错误（缺少冒号、行级后续列是否含1的逻辑未正确实现）
• 未针对每行做列后是否有1的判断，仅做了列级判断

正确实现代码 #

import numpy as np
import pandas as pd

n = 4
t = 3   
condition = pd.DataFrame([[0,0,1,0], [0,1,0,0], [0,0,0, 1], [0,0,0, 1]], columns=['0','1', '2', '3'])
condition.index = np.arange(1, n+1)

# 初始化DataFrames
df = pd.DataFrame(index=range(1, n+1), columns=range(t+1))
df_2 = pd.DataFrame(index=range(1, n+1), columns=range(t+1))
df_3 = pd.DataFrame(3, index=range(1, n+1), columns=range(t+1))

# 先获取每行最后一个1所在的列索引，快速判断后续列是否有1
last_one_col = condition.eq(1).idxmax(axis=1)
# 注：本次数据所有行都有1，无需处理无1的情况

# 从最后一列往前循环
for col in range(t, -1, -1):
    for idx in df.index:
        if condition.loc[idx, str(col)] == 1:
            # 条件1：condition值为1，填充df_3的平方
            df.loc[idx, col] = df_3.loc[idx, col] ** 2
            df_2.loc[idx, col] = 0
        else:
            if col > last_one_col.loc[idx]:
                # 当前列在最后一个1的列之后，后续无1，填充2.5
                df.loc[idx, col] = 2.5
            else:
                # 当前列在最后一个1的列之前，后续有1，填充5
                df.loc[idx, col] = 5
                # 按原逻辑填充df_2，避免最后一列越界
                if col < t:
                    df_2.loc[idx, col] = df.loc[idx, col+1]

# 打印最终结果
print(df)

代码说明 #

1. 先计算每行最后一个1的列索引last_one_col，用于快速判断当前列之后是否存在1
2. 采用列倒序循环（从最后一列到第一列），同时遍历每行，逐元素处理
3. 针对每个元素按需求的三个分支赋值，同时处理df_2的填充逻辑
4. 处理df_2时增加边界判断，避免最后一列访问col+1出现越界错误

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