python中是否有retain函数？我想将SAS代码转移到python

2019年7月2日 338次阅读

我想将SAS代码转移到 python,并且在 python中找不到retain函数.

数据如下：

type_id    amount
1           100
1           200
1           400
2           0
1           200
1           300
2           0
1           150

我想要的是当type_id = 2时,金额等于前一行的负值.所以保留函数后数据会是这样的：

type_id    amount
1           100
1           200
1           400
2          -400
1           200
1           300
2          -300
1           150

SAS代码是：

data B;
set A;
retain tempvar 0;
if type_id = 2
then amount = tempvar;
else tempvar = -amount;
drop tempvar;
run;

有没有人知道如何在python中这样做？
谢谢！

最佳答案 IIUC

df
type_id amount
0   1   100
1   1   200
2   1   400
3   2   0
4   1   200
5   1   300
6   2   0
7   1   150

def retain(df):
    df['ret'] = df['amount'].shift()
    df.ix[df['type_id']==2,'amount'] = -df.ix[df['type_id']==2,'ret']
    df.drop("ret", axis=1, inplace=True)
    return df

retain(df)
type_id amount
0   1   100.0
1   1   200.0
2   1   400.0
3   2   -400.0
4   1   200.0
5   1   300.0
6   2   -300.0
7   1   150.0

或者：

def retain(df):
    df.amount.ix[df.type_id==2] = - df.amount.shift().ix[df.type_id==2]
    return df

retain(df)
type_id amount
0   1   100.0
1   1   200.0
2   1   400.0
3   2   -400.0
4   1   200.0
5   1   300.0
6   2   -300.0
7   1   150.0
点赞