根据现有列中的条件在数据框中创建新列

如何解决根据现有列中的条件在数据框中创建新列

我有以下数据：

dict1={"Code":[3,3,1,2,3],"Num":[10,10,5,25,25]}

df1=pd.DataFrame(dict1)

结果为：

我想创建一个新列（期末余额），其值基于现有的Code和Num列。

如果“代码”值为1，则“余额”等于“数字”

如果代码为2，则“末余额”值为代码为2时Num值的总和。

如果代码为3，则“最终余额”值为“代码”为3时Num值的总和。

我使用迭代，并且具有以下脚本：

mylist1=[]
mylist2=[]
for index,row in df1.iterrows():
    if row["Code"]==1:
        end_balance=row["Num"]  
    elif row["Code"]==2:
        mylist1.append(row["Num"])
        end_balance=sum(mylist1) 
    elif row["Code"]==3:
        mylist2.append(row["Num"])
        end_balance=sum(mylist2)
    df1.loc[index,"End_Balance"]=end_balance

哪个输出是

   Code Num End_Balance
0   3   10  10.00
1   3   10  20.00
2   3   5   25.00
3   1   5   5.00
4   1   5   5.00
5   2   5   5.00
6   2   10  15.00
7   3   5   30.00
8   3   25  55.00
9   3   25  80.00

此输出的问题在于，在Code = 3的第二个子集处，End_Balance列开始求和，同时考虑到Code为3的第一个子集。您可以轻松地看到这一点。我希望以某种方式在Code = 3的第一个子集之后删除脚本中的mylist2，并且当新的Code = 3的子集出现时，End_Balance列中的求和应该重新开始。预期输出为：

   Code Num End_Balance
0   3   10  10.00
1   3   10  20.00
2   3   5   25.00
3   1   5   5.00
4   1   5   5.00
5   2   5   5.00
6   2   10  15.00
7   3   5   5.00
8   3   25  30.00
9   3   25  55.00

希望您的建议遵循相同的逻辑-使用迭代。我知道有可能通过groupby可以做我想做的事，但是在这里我需要一个解决问题的解决方案。

解决方法

IIUC，np.where + GroupBy.cumsum

import numpy as np
blocks = df['Code'].ne(df['Code'].shift()).cumsum()
df['End_balance'] = np.where(df['Code'].eq(1),df['Num'],df.groupby(blocks)['Num'].cumsum())

print(df)

   Code  Num  End_balance
0     3   10           10
1     3   10           20
2     3    5           25
3     1    5            5
4     1    5            5
5     2    5            5
6     2   10           15
7     3    5            5
8     3   25           30
9     3   25           55

或Series.where：

df['End_balance'] = df['Num'].where(df['Code'].eq(1),df.groupby(blocks)['Num'].cumsum())

您也可以使用np.select：

In [2062]: import numpy as np

In [2063]: conditions = [df1.Code.eq(1),df1.Code.eq(2) | df1.Code.eq(3)]
In [2061]: choices = [df1.Num,df1.groupby((df1.Code != df1.Code.shift()).cumsum())['Num'].cumsum()]

In [2065]: df1['End_Balance'] = np.select(conditions,choices)

In [2066]: df1
Out[2066]: 
    Code  Num  End_Balance
0     3   10           10
1     3   10           20
2     3    5           25
3     1    5            5
4     1    5            5
5     2    5            5
6     2   10           15
7     3    5            5
8     3   25           30
9     3   25           55

df1['End_Balance'] = np.where(df1.Code == 1,df1.Num,df1.groupby((df1.Code != df1.Code.shift(1)).cumsum())['Num'].transform('cumsum') )
print(df1)

打印：

   Code  Num  End_Balance
0     3   10           10
1     3   10           20
2     3    5           25
3     1    5            5
4     1    5            5
5     2    5            5
6     2   10           15
7     3    5            5
8     3   25           30
9     3   25           55