python - Pandas:按行扩展数据帧，类似于 R 的 SurvSplit()

Question

我有一个包含人员和工作年限的数据框:

person_id  years                
1          1.00
2          2.34
3          6.85

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

因此，理想情况下，我想提供 block 的 list 或 tuple 来指导按行扩展(例如 [1,2] code> 在第 1 年和第 2 年进行分块)

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

我该怎么做？我在 Stackoverflow 上找到了几篇文章，但它们讨论了不同的数据帧扩展目标。

Answer 1

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

下面基本上运行了迭代任期年数与最大块参数的交叉连接，并合并到人员的年点。然后，将原始数据帧值与计算出的 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