如何按规则重塑Pandas DataFrame：添加结束时间行与空行

需求说明 #

需要对Pandas DataFrame执行以下重塑操作：

• 为每一行新增一行：该行的ts值为原行ts与dur的和，其余列与原行保持一致；
• 在每一组（原行+新增行）之后添加一行全为NA的空行。

现有代码及输出 #

用户当前的实现代码：

import pandas as pd
import itertools

def segments(df):
    start = df['ts']
    end = start + df['dur']
    lst = list(zip(start, end, itertools.cycle([None])))
    return lst

df = pd.DataFrame({
    'ts':   [1, 5, 10 ],
    'dur':  [1, 2, 1 ],
    'typ':  [0, 1, 0]
})
print(df)
lst = segments(df)
print(lst)

当前输出：

ts  dur  typ
0   1    1    0
1   5    2    1
2  10    1    0
[(1, 2, None), (5, 7, None), (10, 11, None)]

期望输出 #

ts  typ
0  1    0
1  2    0
2  NA   NA
3  5    1
4  7    1
5  NA   NA
6  10   0
7  11   0
8  NA   NA

解决方案 #

方法1：遍历行构造（直观易读）

通过遍历原DataFrame的每一行，依次添加原行、计算后的行和空行，最后组合成新DataFrame：

import pandas as pd

def reshape_df(df):
    new_rows = []
    for _, row in df.iterrows():
        # 添加原行（保留ts和typ列）
        new_rows.append({'ts': row['ts'], 'typ': row['typ']})
        # 添加ts=原ts+dur的行
        new_rows.append({'ts': row['ts'] + row['dur'], 'typ': row['typ']})
        # 添加全NA空行
        new_rows.append({'ts': pd.NA, 'typ': pd.NA})
    # 转换为DataFrame并重置索引
    return pd.DataFrame(new_rows).reset_index(drop=True)

# 测试
df = pd.DataFrame({
    'ts':   [1, 5, 10 ],
    'dur':  [1, 2, 1 ],
    'typ':  [0, 1, 0]
})

result = reshape_df(df)
print(result)

方法2：向量化实现（高效适合大数据集）

利用Pandas的向量化操作避免循环，提升处理效率：

import pandas as pd

def reshape_df_vectorized(df):
    # 扩展原行和计算后的行
    expanded = df.assign(ts=df.apply(lambda x: [x['ts'], x['ts']+x['dur']], axis=1)).explode('ts')
    # 生成对应数量的全NA行
    na_df = pd.DataFrame([pd.NA]*len(df), index=expanded.index[::2]+0.5, columns=['typ']).reset_index(drop=True)
    # 合并扩展行和NA行并排序索引
    combined = pd.concat([expanded[['ts','typ']], na_df]).sort_index().reset_index(drop=True)
    # 填充NA行的ts列
    combined['ts'] = combined['ts'].fillna(pd.NA)
    return combined

# 测试
df = pd.DataFrame({
    'ts':   [1, 5, 10 ],
    'dur':  [1, 2, 1 ],
    'typ':  [0, 1, 0]
})

result = reshape_df_vectorized(df)
print(result)

两种方法都能得到符合期望的输出，可根据数据规模选择合适的实现方式。

内容的提问来源于stack exchange，提问作者lucky1928