Pandas向量化操作：实现序列连续计数与指定阈值重置

引言：序列计数与阈值重置的挑战

在数据分析场景中，我们经常需要对序列中连续出现的相同值进行计数。例如，在金融数据中，可能需要统计股票价格连续上涨或下跌的天数。更进一步，有时还需要引入一个重置机制，即当连续计数达到某个预设阈值时，计数器自动归零并重新开始计数。传统的python循环方法虽然可以实现这一逻辑，但在处理大型数据集时效率低下，难以满足高性能要求。pandas作为数据处理的强大工具，提供了丰富的向量化操作，能够以更优雅、更高效的方式解决此类问题。

本教程将详细阐述如何利用Pandas的groupby、shift、cumsum和cumcount等核心功能，结合数学运算，实现一个既能统计连续序列，又能根据指定阈值（例如5）进行重置的计数器。

核心概念与方法

要实现带阈值重置的连续计数，我们需要解决两个关键问题：

1. 识别连续块：如何将DataFrame中连续相同的元素划分为不同的组。
2. 组内计数与阈值重置：如何在每个连续组内进行累积计数，并在达到阈值时重置。

1. 识别连续块 (ne().shift().cumsum())

识别连续块是解决问题的首要步骤。我们可以通过比较当前元素与其前一个元素是否相等来判断连续性。

• df['sign'].shift()：将sign列向下移动一位，使得每一行可以访问到其前一行的sign值。
• df['sign'].ne(df['sign'].shift())：比较当前sign值是否不等于前一个sign值。这将生成一个布尔序列，True表示sign发生了变化（即一个新的连续块开始），False表示sign与前一个相同（仍在同一连续块内）。对于第一行，由于shift()的结果是NaN，ne()会将其视为True，从而确保第一个块被正确识别。
• .cumsum()：对布尔序列进行累积求和。由于True在数值运算中被视为1，False被视为0，cumsum()会为每个新的连续块生成一个唯一的递增整数ID。

例如： | sign | sign.shift() | sign.ne(sign.shift()) | cumsum() (consecutive block ID) | | :--- | :----------- | :-------------------- | :------------------------------ | | 1 | NaN | True | 1 | | 1 | 1 | False | 1 | | -1 | 1 | True | 2 | | -1 | -1 | False | 2 | | 1 | -1 | True | 3 | | 1 | 1 | False | 3 |

2. 组内计数与阈值重置 (groupby().cumcount() % threshold + 1)

一旦我们有了连续块的ID，就可以利用groupby()对这些块进行分组，并在每个组内进行累积计数。

• df.groupby(consecutive_block_id)：根据上一步生成的连续块ID对DataFrame进行分组。
• .cumcount()：对每个组内的元素进行累积计数。它从0开始为每个组内的第一个元素计数，然后递增。
  • 例如，在ID为1的组中，计数将是0, 1, 2...
  • 在ID为2的组中，计数将是0, 1, 2...
• % threshold + 1：这是实现阈值重置的关键。
  • % threshold：取模运算。当cumcount()达到threshold时，cumcount() % threshold的结果将是0。例如，当threshold为5时，cumcount()为0,1,2,3,4,5,6...，则cumcount() % 5为0,1,2,3,4,0,1...。
  • + 1：由于cumcount()和取模运算的结果都是从0开始，为了使计数从1开始（符合人类直觉），我们需要将结果加1。最终，计数序列将变为1,2,3,4,5,1,2...，完美实现了在达到5时重置为1。

完整解决方案

结合上述概念，我们可以用一行Pandas代码实现所需功能：

import pandas as pd

# 示例数据
data = {
    'price': [13, 12, 11, 12, 13, 14, 14, 14, 14, 14, 14, 15, 16],
    'sign': [1, 1, -1, -1, 1, 1, 1, 1, 1, 1, 1, -1, -1]
}
df = pd.DataFrame(data)

# 定义重置阈值
reset_threshold = 5

# 实现连续计数与阈值重置
df['count'] = df.groupby(df['sign'].ne(df['sign'].shift()).cumsum()).cumcount() % reset_threshold + 1

print(df)

详细解析与中间步骤

为了更好地理解上述代码的执行流程，我们可以逐步分解其内部操作，并观察中间结果。

序列猴子开放平台

具有长序列、多模态、单模型、大数据等特点的超大规模语言模型

0

查看详情

import pandas as pd

data = {
    'price': [13, 12, 11, 12, 13, 14, 14, 14, 14, 14, 14, 15, 16],
    'sign': [1, 1, -1, -1, 1, 1, 1, 1, 1, 1, 1, -1, -1]
}
df = pd.DataFrame(data)
reset_threshold = 5

# 步骤1: 识别符号变化点
# df['sign'].shift() 将 'sign' 列向下移动一位
df['sign_shifted'] = df['sign'].shift()
# df['sign'].ne(df['sign_shifted']) 比较当前 'sign' 是否不等于前一个 'sign'
df['is_new_block'] = df['sign'].ne(df['sign_shifted'])

# 步骤2: 生成连续块的唯一ID
# .cumsum() 对布尔值求和，为每个连续块生成一个递增ID
df['consecutive_block_id'] = df['is_new_block'].cumsum()

# 步骤3: 在每个连续块内进行累积计数（从0开始）
# df.groupby(...).cumcount() 对每个组内的元素进行累积计数
df['raw_cumcount'] = df.groupby(df['consecutive_block_id']).cumcount()

# 步骤4: 应用阈值重置并调整为从1开始计数
# % reset_threshold 实现计数重置
# + 1 将计数调整为从1开始
df['final_count'] = df['raw_cumcount'] % reset_threshold + 1

print(df)

输出结果：

    price  sign  sign_shifted  is_new_block  consecutive_block_id  raw_cumcount  final_count
0      13     1           NaN          True                     1             0            1
1      12     1           1.0         False                     1             1            2
2      11    -1           1.0          True                     2             0            1
3      12    -1          -1.0         False                     2             1            2
4      13     1          -1.0          True                     3             0            1
5      14     1           1.0         False                     3             1            2
6      14     1           1.0         False                     3             2            3
7      14     1           1.0         False                     3             3            4
8      14     1           1.0         False                     3             4            5
9      14     1           1.0         False                     3             5            1
10     14     1           1.0         False                     3             6            2
11     15    -1           1.0          True                     4             0            1
12     16    -1          -1.0         False                     4             1            2

从上述详细输出中，我们可以清晰地看到consecutive_block_id如何将连续的sign值分组，raw_cumcount如何在每个组内从0开始计数，以及final_count如何通过取模运算在达到5时重置为1。

最终结果

移除中间辅助列，我们得到期望的输出：

    price  sign  count
0      13     1      1
1      12     1      2
2      11    -1      1
3      12    -1      2
4      13     1      1
5      14     1      2
6      14     1      3
7      14     1      4
8      14     1      5
9      14     1      1
10     14     1      2
11     15    -1      1
12     16    -1      2

注意事项与应用场景

• 性能优势：Pandas的向量化操作在底层由C语言实现，相比Python循环具有显著的性能优势，尤其适用于大规模数据集。
• 可读性：虽然一行代码可能看起来有些复杂，但一旦理解了groupby().cumcount()和shift().cumsum()的组合逻辑，它比等效的循环代码更简洁、更易于维护。
• 灵活性：
  • 修改重置阈值：只需修改reset_threshold变量的值即可。
  • 适用于任何列：这种方法不仅适用于sign列，也可以应用于DataFrame中任何需要对连续相同值进行计数的列。
  • 不同计数起始值：如果希望计数从0开始，可以移除+ 1。
• 处理初始NaN：shift()操作会在第一行产生NaN。ne()（不等于）操作符会将其与df['sign'][0]比较，结果通常是True，这使得cumsum()能正确地将第一行标记为第一个连续块的开始。

总结

本教程展示了如何利用Pandas强大的向量化能力，通过巧妙地组合shift()、ne()、cumsum()、groupby()和cumcount()，并结合取模运算，高效地实现了序列连续计数以及指定阈值重置的功能。这种方法不仅代码简洁、易于维护，更重要的是在处理大量数据时表现出卓越的性能。掌握此类高级Pandas技巧，能够显著提升数据处理的效率和代码质量。

以上就是Pandas向量化操作：实现序列连续计数与指定阈值重置的详细内容，更多请关注php中文网其它相关文章！