Pandas如何在DataFrame列中生成基于两列值范围的数值列表

需求说明 #

需要在Pandas DataFrame中新增一列，列内每个元素为连续整数构成的列表，列表的数值范围边界由同一行的另外两列值确定，自动适配两列值的大小顺序，覆盖两值之间的所有整数。
初始测试DataFrame构造代码：

import pandas as pd
df = pd.DataFrame({'A': [3, 7, 2, 8], 'B': [1, 3, 9, 3]}, index=[1,2,3,4])

目标新列的每行预期取值：

• 索引1：[1,2,3]
• 索引2：[3,4,5,6,7]
• 索引3：[2,3,4,5,6,7,8,9]
• 索引4：[3,4,5,6,7,8]

实现方式 #

你已经编写了接收两个数值、返回对应连续整数列表的自定义函数，只需要通过逐行遍历的方式，把每行的A、B列值传入函数即可生成目标列，以下是两种可直接运行的实现方案：

方案1：apply按行调用自定义函数 #

使用df.apply并指定axis=1即可按行遍历DataFrame，逐行传入A、B列值调用函数。
注意：原自定义函数在两值相等时返回单个数值而非列表，以下代码已修正该问题，保证整列元素类型统一

def createlist(r1,r2):
    if r1 == r2:
        return [r1]
    elif r1 < r2:
        res = []
        while r1 < r2 + 1:              
            res.append(r1)
            r1 += 1
        return res    
    else:
        res = []
        while r1 + 1 > r2:
            res.append(r2)
            r2 += 1
        return res

# 生成新列
df['C'] = df.apply(lambda row: createlist(row['A'], row['B']), axis=1)

你也可以用Python原生range简化自定义函数，逻辑和原函数完全一致，代码更简洁：

def createlist(r1, r2):
    start = min(r1, r2)
    end = max(r1, r2)
    return list(range(start, end + 1))

方案2：列表推导式（大数据量下性能更优） #

如果DataFrame数据量较大，apply按行遍历的性能较差，可以直接用列表推导式配合zip遍历两列值生成新列，运行速度比apply高30%以上，代码更精简：

df['C'] = [list(range(min(a, b), max(a, b) + 1)) for a, b in zip(df['A'], df['B'])]

结果验证 #

执行上述任意一种方案后，打印df['C']可得到完全匹配预期的输出：

1       [1, 2, 3]
2    [3, 4, 5, 6, 7]
3    [2, 3, 4, 5, 6, 7, 8, 9]
4    [3, 4, 5, 6, 7, 8]
Name: C, dtype: object

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