如何基于条件从numpy数组向不同大小的DataFrame分配连续值

需要基于条件，从numpy数组向大小不同的pandas DataFrame列分配连续值。示例代码及当前输出如下：

import pandas as pd
import numpy as np

if __name__ == "__main__":
    df = pd.DataFrame([np.nan, 1, np.nan, 1, np.nan, 1, np.nan])
    arr = np.array([4, 5, 6])

    i = iter(arr)

    df[0] = np.where(df[0] == 1, next(i), np.nan)

    print(df)

当前输出：

0
0  NaN
1  4.0
2  NaN
3  4.0
4  NaN
5  4.0
6  NaN

期望输出：

0
0  NaN
1  4.0
2  NaN
3  5.0
4  NaN
5  6.0
6  NaN

方法1：布尔索引直接赋值 #

通过布尔索引定位DataFrame中值为1的行，直接将numpy数组的元素对应赋值到这些位置：

import pandas as pd
import numpy as np

if __name__ == "__main__":
    df = pd.DataFrame([np.nan, 1, np.nan, 1, np.nan, 1, np.nan])
    arr = np.array([4, 5, 6])
    
    mask = df[0] == 1
    df.loc[mask, 0] = arr
    
    print(df)

方法2：利用replace方法替换 #

使用replace方法将所有值为1的位置按顺序替换为numpy数组的元素：

import pandas as pd
import numpy as np

if __name__ == "__main__":
    df = pd.DataFrame([np.nan, 1, np.nan, 1, np.nan, 1, np.nan])
    arr = np.array([4, 5, 6])
    
    df[0] = df[0].replace(1, iter(arr))
    
    print(df)

方法3：补全数组后赋值 #

先创建和DataFrame同长度的NaN数组，再将对应位置替换为numpy数组的值：

import pandas as pd
import numpy as np

if __name__ == "__main__":
    df = pd.DataFrame([np.nan, 1, np.nan, 1, np.nan, 1, np.nan])
    arr = np.array([4, 5, 6])
    
    fill_arr = np.full(len(df), np.nan)
    fill_arr[df[0] == 1] = arr
    
    df[0] = fill_arr
    
    print(df)

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