Pandas：按行展开dataframe，类似于R的SurvSplit（）

Question

我有一个雇员和工作年限的数据框架：

person_id  years                
1          1.00
2          2.34
3          6.85

我想根据雇员任期中预定义的“块”逐行扩展dataframe。例如，如果我想将人们的任期划分为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

因此，理想情况下，我想提供一个list或tuple的块来指导行扩展（例如[1,2]到第1年和第2年的块）

这种数据帧操作类似于R的survSplit()——参见第127页here

我该怎么做？我发现了一些关于Stackoverflow的文章，但它们讨论了不同的数据帧扩展目标。你知道吗

Answer 1

考虑以下定义的方法。虽然有点走查，但它没有使用循环，不像survsplit实际的源代码是用C编写的

下面基本上运行一个交叉连接迭代任期年到最大块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