如何在DataFrame中生成指定规则的数值范围列表作为新列？

原始DataFrame #

import pandas as pd

df = pd.DataFrame(
        {
           'a': [20, 100],
           'b': [2, 3],
           'dir': ['long', 'short']
        }
 )

预期输出 #

需新增列x，最终DataFrame如下：

a  b    dir    x
0   20  2   long    [22, 24, 26]
1  100  3  short    [97, 94, 91]

规则说明 #

• x为长度3的列表
• a是序列起始值，b为步长
• 根据dir判断增减方向：dir为long时序列升序，为short时降序

尝试的代码（未达成预期） #

df['x'] = np.arange(0, 3) * df.b + df.a

正确解法 #

方法1：使用apply逐行处理 #

import numpy as np

def build_x(row):
    step = row['b'] if row['dir'] == 'long' else -row['b']
    # 生成从a+step开始的3个连续值
    return [row['a'] + step * i for i in range(1, 4)]

df['x'] = df.apply(build_x, axis=1)

方法2：向量化操作（效率更高） #

import numpy as np

# 生成方向系数：long对应1，short对应-1
dir_factor = np.where(df['dir'] == 'long', 1, -1)
# 通过广播生成每行的3个值，再转为列表
df['x'] = (df['a'].to_numpy()[:, None] + dir_factor.to_numpy()[:, None] * df['b'].to_numpy()[:, None] * np.arange(1,4)).tolist()

两种方法均可生成符合要求的x列。

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