如何编写函数移除Pandas DataFrame中col2列的前置零行？

原始数据 #

首先定义目标DataFrame：

import pandas as pd

df = pd.DataFrame({
    'n': [0,1,2,3, 0,1,2, 0,1,2],
    'col1': ['A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'C', 'C'],
    'col2': [0, 0, 0, 0, 3.3, 0, 4, 1.94, 0, 6.17]
})

原始数据展示：

n col1  col2
0   0    A  0.00
1   1    A  0.00
2   2    A  0.00
3   3    B  0.00
4   0    B  3.30
5   1    B  0.00
6   2    B  4.00
7   0    C  1.94
8   1    C  0.00
9   2    C  6.17

需求说明 #

需要编写函数移除col2列中前置的连续零行，仅保留第一个非零行及之后的所有数据（包括后续出现的零行），目标结果如下：

n col1  col2
4   0    B  3.30
5   1    B  0.00
6   2    B  4.00
7   0    C  1.94
8   1    C  0.00
9   2    C  6.17

原函数存在的问题：

def remove_lead_zeros(df):
   new_df = df[df['col2'] != 0]
   return new_df

该函数会删除所有col2值为0的行，不符合仅移除前置连续零行的需求。

解决方案 #

核心逻辑 #

定位col2列中第一个非零值所在的行索引，从该行开始截取DataFrame，即可保留第一个非零行及之后的所有数据（包括后续零行），同时处理全零的边界情况。

实现代码 #

def remove_lead_zeros(df):
    # 筛选出col2非零的行
    non_zero_rows = df[df['col2'] != 0]
    # 若全为零，返回原数据副本
    if non_zero_rows.empty:
        return df.copy()
    # 获取第一个非零行的索引
    first_non_zero_idx = non_zero_rows.index[0]
    # 截取并返回目标数据
    return df.loc[first_non_zero_idx:]

验证结果 #

调用上述函数后，返回的DataFrame与预期目标结果完全一致。

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