通过 pandas 中的函数获取列表

Posted 2023-03-12

【中文标题】通过 pandas 中的函数获取列表

【英文标题】：Getting a list as the result of a function in pandas

【发布时间】：2017-02-03 02:03:31

【问题描述】：

我在 pandas 中有数据框，并且我编写了一个函数来使用每一行中的信息来生成一个新列。我希望结果为列表格式：

      A    B    C
      3    4    1
      4    2    5

     def Computation(row):
          if row['B'] >= 3:
              return [s for s in range(row['C'],50)]
          else:
              return [s for s in range(row['C']+2,50)]

     df['D'] = df.apply(Computation, axis = 1) 

但是，我收到以下错误：

“无法将输入数组从形状 (308) 广播到形状 (9)”

你能告诉我如何解决这个问题吗？

【参考方案1】：

假设你开始

In [25]: df = pd.DataFrame('A': [3, 4], 'B': [4, 2], 'C': [1, 5])

那么至少有两种方法可以做到。

您可以在C 列上应用两次，但打开B 列：

In [26]: np.where(df.B >= 3, df.C.apply(lambda c: [s for s in range(c, 50)]), df.C.apply(lambda c: [s for s in range(c + 2, 50)]))
Out[26]: 
array([ [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
       [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49]], dtype=object)

或者您可以应用整行并打开每行的B 值：

In [27]: df.apply(lambda r: [s for s in range(r.C, 50)] if r.B >= 3 else [s for s in range(r.C + 2, 50)], axis=1)
Out[27]: 
0    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14...
1    [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, ...

请注意，返回类型不同，但在每种情况下，您仍然可以编写

df['foo'] = <each one of the above options>

【讨论】：

这很棒。但是，如果范围的起点取决于每行中 B 的值，我该如何实现呢？

@user36729 你能举一个具体的例子来说明你想做什么吗？ b 和 4c 之间的范围是一个很好的例子吗？

我在主要问题中这样做了。

@user36729 哦，我现在看到您确实切换了问题的内容。查看更新。

对不起，第二个起初对我有用，但现在我收到此错误：“无法将输入数组从形状 (308) 广播到形状 (9)”。