根据条件将迭代器中的值插入pandas列

Question

假设有一个类似的数据帧

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

我想在列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)

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

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。我发现的唯一一个解决方案是使用中间数据帧

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解决方案，而不是创建一个新的临时数据帧？或者有任何pandasbultin函数可以这样做吗？

Answer 1

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

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

首先，为您的示例创建一个数据帧：

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

您可以使用cumsum((cummulative sum来计算您遇到的负数数量，并将其分配给新列

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

A B0-1 11 2 12-1 23 0 24-1 35-1 46 2 47 3 48 5 49-1 5

这仍然会在a列中有一个正数的值，所以你可以用NaN值代替这些值

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

你最终得到：

A B0-1 1.01 2 NaN2-1 2.03 0 NaN4-1 3.05-1 4.06 2 NaN7 3 NaN8 5 NaN9-1 5.0

Answer 2

通过range:对可能的设置值使用DataFrame.loc

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)