无需for循环实现DataFrame末尾连续相同值的计数统计

需求说明 #

判断DataFrame的position列最后一个值是否不为0，若满足则从末尾反向统计连续等于该值的行数，直到值发生变化，最终将计数存入变量，禁止使用for循环，可借助loc/iloc或pandas/numpy向量化操作实现。

示例数据 #

| position |
0          1
1          0  
2          1  
3          1  
4          1  
5          1  

预期计数结果：count = 4

现有循环实现（需替换） #

count = 1
if data['position'].iloc[-1] != 0:
   for i in data['position']:
      if data['position'].iloc[-count] == data['position'].iloc[-1]:
         count = count + 1 
      else:
         break
   if data['position'].iloc[-count] != data['position'].iloc[-1]:
      count = count - 1 

非循环解决方案 #

利用pandas和numpy的向量化操作实现，效率远高于循环（尤其适用于大数据集）：

import pandas as pd
import numpy as np

# 示例DataFrame
data = pd.DataFrame({'position': [1, 0, 1, 1, 1, 1]})

count = 0
last_val = data['position'].iloc[-1]

if last_val != 0:
    # 生成布尔掩码：标记所有等于最后一个值的位置
    mask = data['position'] == last_val
    # 反转掩码数组，找到第一个不匹配的位置（即第一个False的索引）
    reversed_mask = mask[::-1].values
    first_non_match_idx = np.argmin(reversed_mask)
    
    # 处理全匹配的情况：如果所有值都等于last_val，argmin返回0
    count = first_non_match_idx if first_non_match_idx != 0 else len(data)

print(count)  # 输出：4

代码逻辑解释 #

1. 获取最后值并判断：先取出position列的最后一个值，确认其不为0才继续统计
2. 生成布尔掩码：用向量化操作标记所有与最后值相等的行
3. 反向查找第一个不匹配项：反转掩码数组后，np.argmin会返回第一个False的索引（因为True对应1，False对应0，argmin取第一个最小值的位置），这个索引正好是从末尾开始连续匹配的行数
4. 全匹配处理：如果所有行都等于最后值，argmin返回0，此时直接取DataFrame的总行数作为计数

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