Pandas：按行展开数据框，类似于 R 的 SurvSplit()

Question

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

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

所以，理想情况下，我想提供一个块list或tuple组来指导行式扩展（例如[1,2]在第 1 年和第 2 年分块）

此数据帧操作类似于 R 的survSplit() - 请参阅 此处的第 127 页

我怎样才能做到这一点？ 我在 Stackoverflow 上找到了几篇文章，但它们讨论了不同的数据框扩展目标。

Answer 1

考虑以下定义的方法。 尽管有一点演练，但它不使用循环，这与用 C 编写的survsplit实际源代码不同。

下面基本上运行了迭代任期年数到块arg 的最大值的交叉连接，并合并到人的年数。 然后，将具有计算出的tstart和tend列的原始数据帧值连接到merge结果上。 必须为原始数据帧分配一个键，这里是人：

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