基于Pandas DataFrame同类型列的连续布尔值判定实现需求

实现思路 #

• 筛选并排序目标列：从DataFrame中提取符合指定命名格式（如column_{id}）且ID在传入列表中的列，按ID升序排列，自动排除无关列。
• 确定检查列范围：对每个目标列，找到其在排序后列表中的索引：
  • 若指定参数n，取当前列及后续最多n个列（避免超出总列数）
  • 若未指定n，取当前列及所有后续列
• 行级条件判断：对选中的列子集，使用all(axis=1)快速判断每行是否全为1/True，生成新的布尔型或整型列。

完整代码实现 #

import pandas as pd
import numpy as np

def get_df_series(df: pd.DataFrame, columns_ids: list, n: int=None) -> pd.DataFrame:
    # 筛选并排序目标列：仅保留符合命名规则且ID在指定列表中的列，按ID升序排列
    target_columns = [f'column_{i}' for i in sorted(columns_ids) if f'column_{i}' in df.columns]
    total_target_cols = len(target_columns)
    
    for col_name in target_columns:
        # 获取当前列在目标列列表中的索引位置
        current_idx = target_columns.index(col_name)
        # 计算需要检查的列的结束索引
        if n is not None:
            # 包含当前列在内，最多取n+1个列（当前列+后续n个），不超过总列数
            end_idx = min(current_idx + n + 1, total_target_cols)
        else:
            # 取当前列到所有后续列
            end_idx = total_target_cols
        # 提取要检查的列子集
        check_cols = target_columns[current_idx:end_idx]
        # 生成新列：判断每行是否全为1，转为整型（如需布尔型可去掉.astype(int)）
        new_col_name = f'{col_name}_all_true'
        df[new_col_name] = df[check_cols].all(axis=1).astype(int)
    
    return df


def create_dataframe(numbers: list) -> pd.DataFrame:
    df = pd.DataFrame()
    # 创建目标列
    for i in numbers:
        df[f'column_{i}'] = np.random.randint(2, size=20)
    # 添加无关列模拟实际场景
    df['side_column'] = np.random.randint(2, size=20)
    df['right_column'] = np.random.randint(2, size=20)
    return df


if __name__=="__main__":
    numbers = [1, 2, 3, 4, 5, 6, 7, 8]
    df = create_dataframe(numbers=numbers)
    # 示例1：指定n=3，检查当前列及后续3个列是否全为1
    df = get_df_series(df=df, columns_ids=numbers, n=3)
    # 示例2：不指定n，检查当前列及所有后续列是否全为1
    # df = get_df_series(df=df, columns_ids=numbers)
    print(df.head())

关键细节说明 #

• 目标列排序：通过sorted(columns_ids)保证列按ID递增排列，确保“后续列”的逻辑符合预期
• 索引越界处理：用min()函数限制结束索引，避免因n过大导致的列不存在问题
• 高效判断：使用Pandas内置的all(axis=1)方法，比逐行循环效率高，适合处理较大的DataFrame
• 灵活输出：默认生成整型列（0/1），去掉.astype(int)即可得到布尔型列（True/False）

内容的提问来源于stack exchange，提问作者Wicked Gummy Bear