熊猫:逐行扩展数据帧,类似于R的SurvSplit()

时间:2017-02-11 19:41:53

标签: python pandas

我有一个人和年使用的数据框:

person_id  years                
1          1.00
2          2.34
3          6.85

我希望根据员工任期内预定义的“块”逐行扩展数据帧。例如,如果我想在1年的时间里限制人们的任期,那么上述数据框将成为:

person_id  tstart  tend                 
1          0.00    1.00
2          0.00    1.00
2          1.00    2.34
3          0.00    1.00
3          1.00    6.85

如果我想在1年和2年标记处进行分块,原始数据框将变为:

person_id  tstart  tend                 
1          0.00    1.00
2          0.00    1.00
2          1.00    2.00
2          2.00    2.34
3          0.00    1.00
3          1.00    2.00
3          2.00    6.85

所以,理想情况下,我想提供一个listtuple个块来指导行式扩展(例如{1}}在第1年和第2年的块中

此数据框操作类似于R [1,2] - 请参阅第127页here

我该怎么做?我在Stackoverflow上发现了一些文章,但他们讨论了不同的数据框扩展目标。

1 个答案:

答案 0 :(得分:1)

考虑以下定义的方法。虽然有点演练但它不使用循环,而不像survsplit实际源代码那样用C语言编写。

下面基本上运行迭代任期年限到最大 chunk arg的交叉连接,并合并到人员的点上。 。然后,在merge结果上连接具有计算的 tstart tend 列的原始数据帧值。必须将密钥分配给原始数据框,此处为

from io import StringIO
import pandas as pd
import numpy as np

persons = pd.read_table(StringIO("""person_id  years                
1          1.00
2          2.34
3          6.85"""), sep="\s+").assign(key = 1)

def expand_tenure(chunk):
    newpersons = persons.assign(tstart = chunk, tend = persons['years'])
    newpersons.loc[newpersons['tend'] < chunk, 'tstart'] = np.floor(persons['years'])

    df = pd.DataFrame({'tstart': list(range(0, chunk)),
                       'tend': list(range(1, chunk+1)),
                       'key': 1})

    mdf = pd.merge(persons, df, on='key')    
    mdf = mdf[mdf['tend'] <= mdf['years']][['person_id', 'tstart', 'tend']]

    cdf = pd.concat([newpersons[['person_id', 'tstart', 'tend']], mdf])\
                    .sort_values(['person_id', 'tstart'])\
                    .drop_duplicates(['person_id', 'tend']).reset_index(drop=True)

    return cdf

输出 (三次运行)

print(expand_tenure(1))
#    person_id  tstart  tend
# 0          1     0.0  1.00
# 1          2     0.0  1.00
# 2          2     1.0  2.34
# 3          3     0.0  1.00
# 4          3     1.0  6.85

print(expand_tenure(4))
#    person_id  tstart  tend
# 0          1     0.0  1.00
# 1          2     0.0  1.00
# 2          2     1.0  2.00
# 3          2     2.0  2.34
# 4          3     0.0  1.00
# 5          3     1.0  2.00
# 6          3     2.0  3.00
# 7          3     3.0  4.00
# 8          3     4.0  6.85

print(expand_tenure(12))
#     person_id  tstart  tend
# 0           1     0.0  1.00
# 1           2     0.0  1.00
# 2           2     1.0  2.00
# 3           2     2.0  2.34
# 4           3     0.0  1.00
# 5           3     1.0  2.00
# 6           3     2.0  3.00
# 7           3     3.0  4.00
# 8           3     4.0  5.00
# 9           3     5.0  6.00
# 10          3     6.0  6.85