根据条件将值从迭代器插入到 pandas 列

Question

假设有一个像这样的 dataframe

   A
0  -1
1  -1
2  1
3  0
4  -2
5  3
6  0

我想在 A 列A值为负数的每一行中添加一个具有增量值的新列。

   A   B
0  -1  1
1  -1  2
2  1   nan
3  0   nan
4  -2  3
5  3   nan
6  0   nan

这是我的代码

ids = iter(range(1, np.sum((df['A'] < 0).values.ravel()) + 1))
df['B'] = np.where(df['A'] < 0, next(ids), np.nan)

不幸的是，我得到的是

   A   B
0  -1  1
1  -1  1
2  1   nan
3  0   nan
4  -2  1
5  3   nan
6  0   nan

我也试过用发电机

def id_generator(max_id):
    curr = 1
    while curr <= max_id:
        yield curr
        curr += 1

df['B'] = np.where(df['A'] < 0, next(id_generator(np.sum((df['A'] < 0).values.ravel()))), np.nan)

作为以前的解决方案，我得到了这个 dataframe

   A   B
0  -1  1
1  -1  1
2  1   nan
3  0   nan
4  -2  1
5  3   nan
6  0   nan

它似乎为它处理的每一行创建了一个迭代器/生成器，因此 id 始终为1 。 我发现的唯一一种解决方案是使用中间 dataframe

index = df[df['A'] < 0].index
new_df = pd.DataFrame(data=[x + 1 for x in range(len(index))], columns=['B'], index=index)
df = df.join(new_df)

所以我的问题是，有没有办法坚持使用np.where解决方案而不是创建一个新的临时 dataframe？ 或者是否有任何pandas函数可以这样做？

Answer 1

使用DataFrame.loc按range获取可能的设置值：

m = df['A'] < 0
df.loc[m, 'B'] = range(1, m.sum() + 1)
print (df)
   A    B
0 -1  1.0
1 -1  2.0
2  1  NaN
3  0  NaN
4 -2  3.0
5  3  NaN
6  0  NaN

使用numpy.where的解决方案：

m = df['A'] < 0
df['B'] = np.where(m, m.cumsum(), np.nan)

Answer 2

几乎任何事情都可以直接使用 numpy 和 pandas 函数来完成。 尽量避免使用迭代器和生成器

我有一个使用几行的解决方案

首先，为您的示例创建一个 dataframe：

import numpy as np
import pandas as pd
df = pd.DataFrame(data={'A': [-1,2,-1,0,-1,-1,2,3,5,-1]})

您可以使用 cumsum() 累积和来计算您遇到的负数数量并将其分配给新列

df['B'] = (df['A'] < 0).cumsum()

    A   B
0   -1  1
1   2   1
2   -1  2
3   0   2
4   -1  3
5   -1  4
6   2   4
7   3   4
8   5   4
9   -1  5

这仍然会在 A 列中有一个正数的值，因此您可以用 NaN 值替换这些值

df.loc[df['A'] >=0, 'B'] = np.NaN

你最终得到：

    A   B
0   -1  1.0
1   2   NaN
2   -1  2.0
3   0   NaN
4   -1  3.0
5   -1  4.0
6   2   NaN
7   3   NaN
8   5   NaN
9   -1  5.0