主页 / coding / 问题 / 79411167
Accepted
Sun Jar
Asked: 2025-02-04 17:28:01 +0800 CST

Pandas 应用函数将列表返回到新列

  • 772

我有一个熊猫数据框：

import pandas as pd
import numpy as np

np.random.seed(150)
df = pd.DataFrame(np.random.randint(0, 10, size=(10, 2)), columns=['A', 'B'])

我想添加一个新的列“C”，其值是“B”列中每三行的组合列表。因此我使用下面的方法来实现我的需求，但是这种方法在数据量很大的时候会很慢。

>>> df['C'] = [df['B'].iloc[i-2:i+1].tolist() if i >= 2 else None for i in range(len(df))]
>>> df
   A  B          C
0  4  9       None
1  0  2       None
2  4  5  [9, 2, 5]
3  7  9  [2, 5, 9]
4  8  3  [5, 9, 3]
5  8  1  [9, 3, 1]
6  1  4  [3, 1, 4]
7  4  1  [1, 4, 1]
8  1  9  [4, 1, 9]
9  3  7  [1, 9, 7]

当我尝试使用 df.apply 函数时，收到一条错误消息：

df['C'] = df['B'].rolling(window=3).apply(lambda x: list(x), raw=False)

TypeError: must be real number, not list

我记得 pandas apply 似乎不返回列表，所以有没有更好的方法？我在论坛上搜索，但找不到有关 apply 和 return 的合适主题。

python

2 个回答

  1. Best Answer

    您可以numpy使用sliding_window_view：

    from numpy.lib.stride_tricks import sliding_window_view as swv

N = 3
df['C'] = pd.Series(swv(df['B'], N).tolist(), index=df.index[N-1:])

    输出：

       A  B          C
0  4  9        NaN
1  0  2        NaN
2  4  5  [9, 2, 5]
3  7  9  [2, 5, 9]
4  8  3  [5, 9, 3]
5  8  1  [9, 3, 1]
6  1  4  [3, 1, 4]
7  4  1  [1, 4, 1]
8  1  9  [4, 1, 9]
9  3  7  [1, 9, 7]
    • 6

  2. 我想你可以从另一个角度改变你的想法，比如说，不是row-wise滑动column-wise窗口，而是除非你的窗口很大，否则你的代码可能会加速N。

    例如，你可以尝试

    N = 3
nr = len(df)
df['C'] = [None]*(N-1) + np.column_stack([df['B'].iloc[k:nr-N+1+k] for k in range(N)]).tolist()

    你将获得

        A   B   C
0   4   9   None
1   0   2   None
2   4   5   [9, 2, 5]
3   7   9   [2, 5, 9]
4   8   3   [5, 9, 3]
5   8   1   [9, 3, 1]
6   1   4   [3, 1, 4]
7   4   1   [1, 4, 1]
8   1   9   [4, 1, 9]
9   3   7   [1, 9, 7]
    • 2

相关问题