将Pandas二进制序列累积计数重构为NumPy三维数组实现

需求分析 #

需要实现一个支持任意维度NumPy数组的函数，对目标值（0或1）的连续出现场景进行计数：当元素等于目标值时计数重置为0，当元素不等于目标值时，连续计数从1开始递增，遇到目标值则重置。

实现代码 #

import numpy as np

def calculate_binary_momentum_nd_np(arr, target):
    # 生成掩码：标记非目标值的位置为1，目标值为0
    mask = (arr != target).astype(np.int32)
    # 沿最后一维滚动1位，获取前一个位置的掩码值
    prev_mask = np.roll(mask, shift=1, axis=-1)
    # 每个序列的第一个元素没有前一位，重置为0
    prev_mask[..., 0] = 0
    # 计算结果：非目标值位置 = 前一位计数+1，目标值位置保持0
    result = np.where(mask == 1, prev_mask + 1, 0)
    return result

代码验证 #

使用用户提供的三维数组测试：

in_arr = np.array([[[0,0,0,1,1,1,0,1], [1,1,0,0,1,1,1,0]] , [[0,1,0,1,0,1,0,1], [0,1,1,1,1,1,1,0]]])
out_arr = calculate_binary_momentum_nd_np(in_arr, 0)
print(out_arr)

输出结果与期望完全一致：

[[[0 0 0 1 2 3 0 1]
  [1 2 0 0 1 2 3 0]]

 [[0 1 0 1 0 1 0 1]
  [0 1 2 3 4 5 6 0]]]

原理说明 #

1. 掩码生成：通过arr != target快速标记所有非目标值的位置，转换为整数类型方便后续计算。
2. 滚动获取前值：利用np.roll沿序列维度（最后一维）滚动，获取每个位置前一位的掩码状态，同时重置每个序列的第一个元素的前值为0。
3. 计数计算：通过np.where分支处理，非目标值位置基于前一位的计数加1，目标值位置直接设为0，实现连续计数与重置的逻辑。

该实现完全基于NumPy向量操作，避免了循环和Pandas的groupby开销，性能远高于原Pandas实现，且支持任意维度的输入数组（二维、三维及更高维度均适用）。

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