如何实现用随机高度块填充DataFrame中的NaN值？

现有如下5x5的DataFrame：

A      B     C     D     E
0   0.0   0.0   0.0   0.0   0.0
1   0.0   0.0   NaN   0.0   NaN
2   0.0   0.0   NaN   0.0   NaN
3   0.0   NaN   0.0   NaN   NaN
4   0.0   NaN   0.0   NaN   NaN

需创建一个函数，接收上述DataFrame、num_blocks、min_height=1、max_height=5作为参数，以随机高度的块填充各列中的NaN值，步骤如下：

• 统计含NaN值的列数：
  • 若该数值小于num_blocks，抛出错误提示增加num_blocks数值
• 检查num_blocks是否大于DataFrame中NaN的总数：
  • 若大于，抛出错误提示减少num_blocks数值
• 否则，以数值递增的块填充NaN值，块高度需在用户指定的min_height和max_height范围内，切换列时数值递增。

示例1：num_blocks=4, min_height=1, max_height=5 #

填充后DataFrame如下（基础情况）：

A      B     C     D     E
0   0.0   0.0   0.0   0.0   0.0
1   0.0   0.0   2     0.0   4
2   0.0   0.0   2     0.0   4
3   0.0   1     0.0   3     4
4   0.0   1     0.0   3     4

示例2：num_blocks=5, min_height=1, max_height=5 #

填充后可参考如下样式：

A      B     C     D     E
0   0.0   0.0   0.0   0.0   0.0
1   0.0   0.0   2     0.0   5
2   0.0   0.0   2     0.0   5
3   0.0   1     0.0   3     5
4   0.0   1     0.0   4     5

或如下样式：

A      B     C     D     E
0   0.0   0.0   0.0   0.0   0.0
1   0.0   0.0   2     0.0   4
2   0.0   0.0   2     0.0   5
3   0.0   1     0.0   3     5
4   0.0   1     0.0   3     5

示例3：num_blocks=8, min_height=1, max_height=5 #

填充后可参考如下样式：

A      B     C     D     E
0   0.0   0.0   0.0   0.0   0.0
1   0.0   0.0   3     0.0   7
2   0.0   0.0   4     0.0   7
3   0.0   1     0.0   5     8
4   0.0   2     0.0   6     8

示例4：num_blocks=7, min_height=2, max_height=4（含更多NaN的DataFrame） #

原始DataFrame：

A      B     C     D     E
0   NaN   0.0   0.0   0.0   0.0
1   NaN   0.0   NaN   0.0   NaN
2   NaN   NaN   NaN   0.0   NaN
3   NaN   NaN   NaN   NaN   NaN
4   0.0   NaN   0.0   NaN   NaN

填充后结果：

A     B    C     D     E
0   1    0.0   0.0   0.0   0.0
1   1    0.0   4     0.0   6
2   2    3     4     0.0   6
3   2    3     4     5     7
4   0.0  3     0.0   5     7

import pandas as pd

def fill_nans(df, num_blocks):
    # 获取包含NaN的列
    nan_columns = df.columns[df.isna().any()].tolist()

    fill_value = 1

    # 遍历含NaN的列，递增填充值
    for col in nan_columns:
        # 获取当前列中NaN的索引
        nan_idx = df[df[col].isna()].index

        for idx in nan_idx:
            df.at[idx, col] = fill_value
            fill_value += 1
            # 若填充值超过num_blocks则重置
            if fill_value > num_blocks:
                fill_value = 1

    return df

# 示例DataFrame
data = {
    'A': [0.0, 0.0, 0.0, 0.0, 0.0],
    'B': [0.0, 0.0, 0.0, None, None],
    'C': [0.0, None, None, 0.0, 0.0],
    'D': [0.0, 0.0, 0.0, None, None],
    'E': [0.0, None, None, None, None]
}

df = pd.DataFrame(data)

print("原始DataFrame:")
print(df)

print("\n填充NaN后的DataFrame:")
df_filled = fill_nans(df, 8)
print(df_filled)

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