根據現有列中的條件在 dataframe 中創建新列

Question

我有以下數據：

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

df1=pd.DataFrame(dict1)

結果是：

   Code Num
0   3   10
1   3   10
2   3   5
3   1   5
4   1   5
5   2   5
6   2   10
7   3   5
8   3   25
9   3   25

我想創建一個新列（期末余額），其值基於現有的代碼和數字列。

如果 Code 值為 1，則 End Balance 等於 Num

如果 Code 為 2，則 End Balance 值為代碼為 2 的 Num 值的總和

如果 Code 為 3，則 End Balance 值為代碼為 3 的 Num 值的總和

我使用 iterrows 並且有以下腳本：

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

哪個 output 是

   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

這個 output 的問題是，在 Code = 3 的第二個子集中，End_Balance 列開始求和，同時考慮到 Code 為 3 的第一個子集。您可以很容易地看到這一點。 我想以某種方式在 Code=3 的第一個子集之后刪除腳本中的 mylist2，並且當 Code=3 的新子集出現時，End_Balance 列中的總和應該重新開始。 預計 output 是：

   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

願你的建議遵循同樣的邏輯——使用 iterrows。 我知道可能通過 groupby 我可以做我想做的事，但在這里我需要一個帶有 iterrows 的解決方案。

Answer 1

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())

Answer 2

您還可以使用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

Answer 3

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