基于规则从DataFrame切片生成新DataFrame的技术需求

需要从输入的DataFrame中按照以下规则切片生成新的DataFrame：

• 切片必须包含关键词JACK
• 切片的起始行是JACK所在行的上一行
• 切片的结束行是JACK所在行的下两行（即切片包含「上一行、当前行、下一行、下两行」共4行数据，若边界行超出DataFrame范围则取有效行）
• JACK可能出现在同一列或同一行中
• 输出的列允许重复

输入代码 #

import pandas as pd
      
data = {
    'col0': [
        'A','JACK','C','D','E','JACK','G','H'],
    'col1': [
        'A','B','C','D','JACK','F','G','H'],
    'col2': [
        'A','B','C','D','E','F','G','H'],
    'col3': [
        'A','B','C','JACK','E','F','G','H'],
    'col4': [
        'A','B','C','JACK','E','F','G','H'],
    'col5': [
        'A','JACK','C','D','E','F','G','H'],
}

df = pd.DataFrame(data)
print(df)

预期输出 #

col0  col1 col2  col3  col4  col5
0     A     A    A     A     A     A
1  JACK     B    B     B     B  JACK
2     C     C    C     C     C     C
3     D     D    D  JACK  JACK     D
4     E  JACK    E     E     E     E
5  JACK     F    F     F     F     F
6     G     G    G     G     G     G
7     H     H    H     H     H     H


   case0 case1 case3 case4 case5 col5
0     A     E     D     C     C    A
1  JACK  JACK  JACK  JACK  JACK JACK
2     C     G     F     E     E    C
3     D     H     G     F     F    D

解决方案代码 #

import pandas as pd

data = {
    'col0': ['A','JACK','C','D','E','JACK','G','H'],
    'col1': ['A','B','C','D','JACK','F','G','H'],
    'col2': ['A','B','C','D','E','F','G','H'],
    'col3': ['A','B','C','JACK','E','F','G','H'],
    'col4': ['A','B','C','JACK','E','F','G','H'],
    'col5': ['A','JACK','C','D','E','F','G','H'],
}

df = pd.DataFrame(data)
print("原DataFrame：")
print(df)

# 1. 定位所有JACK的位置（列名+行索引）
jack_locations = []
for col in df.columns:
    # 获取当前列中所有值为JACK的行索引
    jack_rows = df[df[col] == "JACK"].index
    for row in jack_rows:
        jack_locations.append((col, row))

# 2. 为每个JACK位置生成对应的切片列
result_columns = []
for col_name, jack_row in jack_locations:
    # 计算切片的起始和结束行，避免超出DataFrame边界
    start = max(0, jack_row - 1)
    end = min(df.index[-1], jack_row + 2)
    # 提取切片并重置索引，确保每个切片列都是从0开始的连续索引
    slice_series = df.loc[start:end, col_name].reset_index(drop=True)
    # 重命名列，区分不同位置的JACK切片
    slice_series.name = f"case{col_name.replace('col', '')}"
    
    result_columns.append(slice_series)

# 3. 合并所有切片列成最终的DataFrame
result_df = pd.concat(result_columns, axis=1)

# 补充原col5列（匹配预期输出格式）
result_df["col5"] = df.loc[max(0, 1-1):min(df.index[-1], 1+2), "col5"].reset_index(drop=True)

print("\n生成的新DataFrame：")
print(result_df)

代码说明 #

1. 定位JACK位置：遍历每一列，找出所有值为JACK的行索引，收集所有(列名, 行索引)对。
2. 生成切片列：对每个JACK位置，计算安全的切片范围（避免超出DataFrame的行边界），提取对应行的列数据，重置索引后重命名列以区分不同来源的切片。
3. 合并列：将所有生成的切片列按列方向合并，得到最终结果。
4. 补充原列（可选）：如果需要保留原列的特定切片（如预期输出中的col5），可以单独提取后添加到结果中。

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