创建数字范围作为数据框的列

Question

我的数据框是：

import pandas as pd

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

预期输出：创建列

x

：

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

步骤：

x

是长度为3的列表。

a

是

x

的起点，

b

是

a

根据

dir

增加/减少的步骤。如果

df.dir == long

x

上升，否则下降。

我的尝试基于此答案：

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

这不会产生预期的输出。

Answer 1

一个可能的选择：

N = 3

sign = {"long": +1, "short": -1}

df["x"] = [
    np.arange(a+sb, a+sb*(N+1), step=sb)# .tolist() ?
    for a,sb in zip(df["a"], df["b"].mul(df["dir"].map(sign)))
]

输出：

print(df)

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

[2 rows x 4 columns]